chore: Fix Typing Issues

core/graphiti.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 import os
+import typing
 from datetime import datetime
 from time import time
 from typing import Callable
@@ -83,7 +84,7 @@
 		episode: EpisodicNode,
 		new_nodes: list[EntityNode],
 		new_edges: list[EntityEdge],
-		relevant_schema: dict[str, any],
+		relevant_schema: dict[str, typing.Any],
 		previous_episodes: list[EpisodicNode],
 	): ...
 
@@ -276,8 +277,8 @@
 			await asyncio.gather(*[node.save(self.driver) for node in nodes])
 
 			# re-map edge pointers so that they don't point to discard dupe nodes
 			extracted_edges: list[EntityEdge] = resolve_edge_pointers(extracted_edges, uuid_map)
 			episodic_edges: list[EpisodicEdge] = resolve_edge_pointers(episodic_edges, uuid_map)
 
 			# save episodic edges to KG
 			await asyncio.gather(*[edge.save(self.driver) for edge in episodic_edges])

core/llm_client/client.py

@@ -1,5 +1,7 @@
+import typing
 from abc import ABC, abstractmethod
 
+from ..prompts.models import Message
 from .config import LLMConfig
 
 
@@ -9,5 +11,5 @@ def __init__(self, config: LLMConfig):
 		pass
 
 	@abstractmethod
-	async def generate_response(self, messages: list[dict[str, str]]) -> dict[str, any]:
+	async def generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
 		pass

core/llm_client/openai_client.py

@@ -1,8 +1,10 @@
 import json
 import logging
+import typing
 
 from openai import AsyncOpenAI
 
+from ..prompts.models import Message
 from .client import LLMClient
 from .config import LLMConfig
 
@@ -14,9 +16,9 @@
 		self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
 		self.model = config.model
 
-	async def generate_response(self, messages: list[dict[str, str]]) -> dict[str, any]:
+	async def generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
 		try:
 			response = await self.client.chat.completions.create(
 				model=self.model,
 				messages=messages,
 				temperature=0.1,

core/nodes.py

@@ -35,8 +35,9 @@ class EpisodicNode(Node):
 	source: str = Field(description='source type')
 	source_description: str = Field(description='description of the data source')
 	content: str = Field(description='raw episode data')
-	valid_at: datetime = Field(
+	valid_at: datetime | None = Field(
 		description='datetime of when the original document was created',
+		default=None,
 	)
 	entity_edges: list[str] = Field(
 		description='list of entity edges referenced in this episode',

core/prompts/dedupe_edges.py

@@ -1,5 +1,5 @@
 import json
-from typing import Protocol, TypedDict
+from typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion
 
@@ -15,7 +15,7 @@ class Versions(TypedDict):
 	edge_list: PromptFunction
 
 
-def v1(context: dict[str, any]) -> list[Message]:
+def v1(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -34,7 +34,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 
         Task:
         1. start with the list of edges from New Edges
-        2. If any edge in New Edges is a duplicate of an edge in Existing Edges, replace the new edge with the existing
+        2. If Any edge in New Edges is a duplicate of an edge in Existing Edges, replace the new edge with the existing
             edge in the list
         3. Respond with the resulting list of edges
 
@@ -55,7 +55,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def v2(context: dict[str, any]) -> list[Message]:
+def v2(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -74,7 +74,7 @@ def v2(context: dict[str, any]) -> list[Message]:
 
         Task:
         1. start with the list of edges from New Edges
-        2. If any edge in New Edges is a duplicate of an edge in Existing Edges, replace the new edge with the existing
+        2. If Any edge in New Edges is a duplicate of an edge in Existing Edges, replace the new edge with the existing
             edge in the list
         3. Respond with the resulting list of edges
 
@@ -97,7 +97,7 @@ def v2(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def edge_list(context: dict[str, any]) -> list[Message]:
+def edge_list(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -112,7 +112,7 @@ def edge_list(context: dict[str, any]) -> list[Message]:
         {json.dumps(context['edges'], indent=2)}
 
         Task:
-        If any edge in Edges is a duplicate of another edge, return the fact of only one of the duplicate edges
+        If Any edge in Edges is a duplicate of another edge, return the fact of only one of the duplicate edges
 
         Guidelines:
         1. Use both the name and fact of edges to determine if they are duplicates, 

core/prompts/dedupe_nodes.py

@@ -1,5 +1,5 @@
 import json
-from typing import Protocol, TypedDict
+from typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion
 
@@ -16,7 +16,7 @@ class Versions(TypedDict):
 	node_list: PromptVersion
 
 
-def v1(context: dict[str, any]) -> list[Message]:
+def v1(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -35,7 +35,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 
         Task:
         1. start with the list of nodes from New Nodes
-        2. If any node in New Nodes is a duplicate of a node in Existing Nodes, replace the new node with the existing
+        2. If Any node in New Nodes is a duplicate of a node in Existing Nodes, replace the new node with the existing
             node in the list
         3. Respond with the resulting list of nodes
 
@@ -56,7 +56,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def v2(context: dict[str, any]) -> list[Message]:
+def v2(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -74,7 +74,7 @@ def v2(context: dict[str, any]) -> list[Message]:
         {json.dumps(context['extracted_nodes'], indent=2)}
 
         Task:
-        If any node in New Nodes is a duplicate of a node in Existing Nodes, add their names to the output list
+        If Any node in New Nodes is a duplicate of a node in Existing Nodes, add their names to the output list
 
         Guidelines:
         1. Use both the name and summary of nodes to determine if they are duplicates, 
@@ -96,7 +96,7 @@ def v2(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def node_list(context: dict[str, any]) -> list[Message]:
+def node_list(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',

core/prompts/extract_edges.py

@@ -1,5 +1,5 @@
 import json
-from typing import Protocol, TypedDict
+from typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion
 
@@ -14,7 +14,7 @@ class Versions(TypedDict):
 	v2: PromptFunction
 
 
-def v1(context: dict[str, any]) -> list[Message]:
+def v1(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -70,7 +70,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def v2(context: dict[str, any]) -> list[Message]:
+def v2(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',

core/prompts/extract_nodes.py

@@ -1,5 +1,5 @@
 import json
-from typing import Protocol, TypedDict
+from typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion
 
@@ -16,7 +16,7 @@ class Versions(TypedDict):
 	v3: PromptFunction
 
 
-def v1(context: dict[str, any]) -> list[Message]:
+def v1(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -64,7 +64,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def v2(context: dict[str, any]) -> list[Message]:
+def v2(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -105,7 +105,7 @@ def v2(context: dict[str, any]) -> list[Message]:
 	]
 
 
-def v3(context: dict[str, any]) -> list[Message]:
+def v3(context: dict[str, Any]) -> list[Message]:
 	sys_prompt = """You are an AI assistant that extracts entity nodes from conversational text. Your primary task is to identify and extract the speaker and other significant entities mentioned in the conversation."""
 
 	user_prompt = f"""

core/prompts/invalidate_edges.py

@@ -1,4 +1,4 @@
-from typing import Protocol, TypedDict
+from typing import Any, Protocol, TypedDict
 
 from .models import Message, PromptFunction, PromptVersion
 
@@ -11,7 +11,7 @@ class Versions(TypedDict):
 	v1: PromptFunction
 
 
-def v1(context: dict[str, any]) -> list[Message]:
+def v1(context: dict[str, Any]) -> list[Message]:
 	return [
 		Message(
 			role='system',
@@ -20,7 +20,7 @@ def v1(context: dict[str, any]) -> list[Message]:
 		Message(
 			role='user',
 			content=f"""
-               Based on the provided existing edges and new edges with their timestamps, determine which existing relationships, if any, should be invalidated due to contradictions or updates in the new edges.
+               Based on the provided existing edges and new edges with their timestamps, determine which existing relationships, if Any, should be invalidated due to contradictions or updates in the new edges.
                 Only mark a relationship as invalid if there is clear evidence from new edges that the relationship is no longer true.
                 Do not invalidate relationships merely because they weren't mentioned in new edges. You may use the current episode and previous episodes as well as the facts of each edge to understand the context of the relationships.
 

core/prompts/lib.py

@@ -1,4 +1,4 @@
-from typing import Protocol, TypedDict
+from typing import Any, Protocol, TypedDict
 
 from .dedupe_edges import (
 	Prompt as DedupeEdgesPrompt,
@@ -68,7 +68,7 @@
 	def __init__(self, func: PromptFunction):
 		self.func = func
 
-	def __call__(self, context: dict[str, any]) -> list[Message]:
+	def __call__(self, context: dict[str, Any]) -> list[Message]:
 		return self.func(context)
 
 
@@ -81,7 +81,7 @@
 class PromptLibraryWrapper:
 	def __init__(self, library: PromptLibraryImpl):
 		for prompt_type, versions in library.items():
 			setattr(self, prompt_type, PromptTypeWrapper(versions))
 
 
 PROMPT_LIBRARY_IMPL: PromptLibraryImpl = {
@@ -92,4 +92,4 @@
 	'invalidate_edges': invalidate_edges_versions,
 }
 
 prompt_library: PromptLibrary = PromptLibraryWrapper(PROMPT_LIBRARY_IMPL)

core/prompts/models.py

@@ -1,4 +1,4 @@
-from typing import Callable, Protocol
+from typing import Any, Callable, Protocol
 
 from pydantic import BaseModel
 
@@ -9,7 +9,7 @@ class Message(BaseModel):
 
 
 class PromptVersion(Protocol):
-	def __call__(self, context: dict[str, any]) -> list[Message]: ...
+	def __call__(self, context: dict[str, Any]) -> list[Message]: ...
 
 
-PromptFunction = Callable[[dict[str, any]], list[Message]]
+PromptFunction = Callable[[dict[str, Any]], list[Message]]

core/utils/maintenance/edge_operations.py

@@ -2,7 +2,7 @@
 import logging
 from datetime import datetime
 from time import time
-from typing import List
+from typing import Any, List
 
 from core.edges import EntityEdge, EpisodicEdge
 from core.llm_client import LLMClient
@@ -35,7 +35,7 @@ async def extract_new_edges(
 	llm_client: LLMClient,
 	episode: EpisodicNode,
 	new_nodes: list[EntityNode],
-	relevant_schema: dict[str, any],
+	relevant_schema: dict[str, Any],
 	previous_episodes: list[EpisodicNode],
 ) -> tuple[list[EntityEdge], list[EntityNode]]:
 	# Prepare context for LLM

core/utils/maintenance/graph_data_operations.py

@@ -52,9 +52,7 @@ async def build_indices_and_constraints(driver: AsyncDriver):
         }}
         """,
 	]
-	index_queries: list[LiteralString] = (
-		range_indices + fulltext_indices + vector_indices
-	)
+	index_queries: list[LiteralString] = range_indices + fulltext_indices + vector_indices
 
 	await asyncio.gather(*[driver.execute_query(query) for query in index_queries])
 
@@ -72,7 +70,6 @@ async def retrieve_episodes(
 	driver: AsyncDriver,
 	reference_time: datetime,
 	last_n: int = EPISODE_WINDOW_LEN,
-	sources: list[str] | None = 'messages',
 ) -> list[EpisodicNode]:
 	"""Retrieve the last n episodic nodes from the graph"""
 	result = await driver.execute_query(

core/utils/maintenance/node_operations.py

@@ -1,4 +1,5 @@
 import logging
+import typing
 from datetime import datetime
 from time import time
 
@@ -12,7 +13,7 @@
 async def extract_new_nodes(
 	llm_client: LLMClient,
 	episode: EpisodicNode,
-	relevant_schema: dict[str, any],
+	relevant_schema: dict[str, typing.Any],
 	previous_episodes: list[EpisodicNode],
 ) -> list[EntityNode]:
 	# Prepare context for LLM

tests/utils/maintenance/test_temporal_operations_int.py

@@ -267,7 +267,7 @@ async def test_invalidate_edges_multiple_invalidations():
 async def test_invalidate_edges_no_effect():
 	existing_edges, nodes = create_complex_test_data()
 
-	# Create a new edge that doesn't invalidate any existing edges
+	# Create a new edge that doesn't invalidate Any existing edges
 	new_edge = (
 		nodes[2],
 		EntityEdge(