working

Author: truemagic-coder

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "solana-agent"
-version = "29.1.3"
+version = "29.2.0-dev80"
 description = "AI Agents for Solana"
 authors = ["Bevan Hunt <bevan@bevanhunt.com>"]
 license = "MIT"
@@ -24,18 +24,18 @@ python_paths = [".", "tests"]
 
 [tool.poetry.dependencies]
 python = ">=3.12,<4.0"
-openai = "1.78.1"
+openai = "1.79.0"
 pydantic = ">=2"
-pymongo = "4.12.1"
-zep-cloud = "2.12.1"
-instructor = "1.8.1"
+pymongo = "4.13.0"
+zep-cloud = "2.12.3"
+instructor = "1.8.2"
 pinecone = "6.0.2"
-llama-index-core = "0.12.35"
+llama-index-core = "0.12.37"
 llama-index-embeddings-openai = "0.3.1"
 pypdf = "5.5.0"
 scrubadub = "2.0.1"
-logfire = "3.15.1"
-typer = "0.15.3"
+logfire = "3.16.0"
+typer = "0.15.4"
 rich = ">=13,<14.0"
 pillow = "11.2.1"
 
@@ -50,7 +50,7 @@ sphinx-rtd-theme = "^3.0.2"
 myst-parser = "^4.0.1"
 sphinx-autobuild = "^2024.10.3"
 mongomock = "^4.3.0"
-ruff = "^0.11.9"
+ruff = "^0.11.10"
 
 [tool.poetry.scripts]
 solana-agent = "solana_agent.cli:app"

pyproject.toml

@@ -0,0 +1,67 @@
+import uuid
+from typing import Dict, Any, List, Optional
+from solana_agent.interfaces.providers.graph_storage import GraphStorageProvider
+from solana_agent.adapters.mongodb_adapter import MongoDBAdapter
+
+
+class MongoDBGraphAdapter(GraphStorageProvider):
+    def __init__(
+        self,
+        mongo_adapter: MongoDBAdapter,
+        node_collection: str = "graph_nodes",
+        edge_collection: str = "graph_edges",
+    ):
+        self.mongo = mongo_adapter
+        self.node_collection = node_collection
+        self.edge_collection = edge_collection
+
+    async def add_node(self, node: Dict[str, Any]) -> str:
+        node = dict(node)
+        node["uuid"] = node.get("uuid", str(uuid.uuid4()))
+        self.mongo.insert_one(self.node_collection, node)
+        return node["uuid"]
+
+    async def add_edge(self, edge: Dict[str, Any]) -> str:
+        edge = dict(edge)
+        edge["uuid"] = edge.get("uuid", str(uuid.uuid4()))
+        return self.mongo.insert_one(self.edge_collection, edge)
+
+    async def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
+        return self.mongo.find_one(self.node_collection, {"uuid": node_id})
+
+    async def get_edges(
+        self, node_id: str, direction: str = "both"
+    ) -> List[Dict[str, Any]]:
+        if direction == "out":
+            query = {"source": node_id}
+        elif direction == "in":
+            query = {"target": node_id}
+        else:
+            query = {"$or": [{"source": node_id}, {"target": node_id}]}
+        return self.mongo.find(self.edge_collection, query)
+
+    async def find_neighbors(
+        self, node_id: str, depth: int = 1
+    ) -> List[Dict[str, Any]]:
+        neighbors = set()
+        current = {node_id}
+        for _ in range(depth):
+            edges = await self.get_edges(list(current)[0])
+            for edge in edges:
+                neighbors.add(edge.get("source"))
+                neighbors.add(edge.get("target"))
+            current = neighbors
+        neighbors.discard(node_id)
+        return [await self.get_node(nid) for nid in neighbors if nid]
+
+    async def temporal_query(
+        self, node_id: str, start_time: Optional[str], end_time: Optional[str]
+    ) -> List[Dict[str, Any]]:
+        query = {"$or": [{"source": node_id}, {"target": node_id}]}
+        if start_time or end_time:
+            query["timestamp"] = {}
+            if start_time:
+                query["timestamp"]["$gte"] = start_time
+            if end_time:
+                query["timestamp"]["$lte"] = end_time
+        return self.mongo.find(self.edge_collection, query)

solana_agent/adapters/mongodb_graph_adapter.py

@@ -10,11 +10,13 @@
 from typing import Dict, Any, List
 
 # Service imports
+from solana_agent.adapters.mongodb_graph_adapter import MongoDBGraphAdapter
 from solana_agent.adapters.pinecone_adapter import PineconeAdapter
 from solana_agent.interfaces.guardrails.guardrails import (
     InputGuardrail,
     OutputGuardrail,
 )
+from solana_agent.services.graph_memory import GraphMemoryService
 from solana_agent.services.query import QueryService
 from solana_agent.services.agent import AgentService
 from solana_agent.services.routing import RoutingService
@@ -241,6 +243,7 @@ def create_from_config(config: Dict[str, Any]) -> QueryService:
 
         # Initialize Knowledge Base if configured
         knowledge_base = None
+        graph_memory_service = None
         kb_config = config.get("knowledge_base")
         # Requires both KB config section and MongoDB adapter
         if kb_config and db_adapter:
@@ -311,6 +314,34 @@ def create_from_config(config: Dict[str, Any]) -> QueryService:
                     "Knowledge Base Service initialized successfully."
                 )  # Use logger.info
 
+                # Create Graph Memory Service
+                graph_memory_config = pinecone_config.get("agent_memory", {})
+                try:
+                    # Create MongoDBGraphAdapter
+                    mongo_graph_adapter = MongoDBGraphAdapter(
+                        mongo_adapter=db_adapter,
+                        node_collection=graph_memory_config.get(
+                            "node_collection", "graph_nodes"
+                        ),
+                        edge_collection=graph_memory_config.get(
+                            "edge_collection", "graph_edges"
+                        ),
+                    )
+
+                    # Create GraphMemoryService
+                    graph_memory_service = GraphMemoryService(
+                        graph_adapter=mongo_graph_adapter,
+                        pinecone_adapter=pinecone_adapter,
+                        openai_adapter=llm_adapter,
+                        embedding_model=graph_memory_config.get(
+                            "embedding_model", "text-embedding-3-large"
+                        ),
+                    )
+                    logger.info("Graph Memory Service initialized successfully.")
+                except Exception as e:
+                    logger.exception(f"Failed to initialize Graph Memory: {e}")
+                    graph_memory_service = None
+
             except Exception as e:
                 # Use logger.exception to include traceback automatically
                 logger.exception(f"Failed to initialize Knowledge Base: {e}")
@@ -324,6 +355,7 @@ def create_from_config(config: Dict[str, Any]) -> QueryService:
             knowledge_base=knowledge_base,  # Pass the potentially created KB
             kb_results_count=kb_config.get("results_count", 3) if kb_config else 3,
             input_guardrails=input_guardrails,
+            graph_memory=graph_memory_service,
         )
 
         return query_service

solana_agent/factories/agent_factory.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Optional
+
+
+class GraphStorageProvider(ABC):
+    @abstractmethod
+    async def add_node(self, node: Dict[str, Any]) -> str:
+        pass
+
+    @abstractmethod
+    async def add_edge(self, edge: Dict[str, Any]) -> str:
+        pass
+
+    @abstractmethod
+    async def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    async def get_edges(
+        self, node_id: str, direction: str = "both"
+    ) -> List[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    async def find_neighbors(
+        self, node_id: str, depth: int = 1
+    ) -> List[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    async def temporal_query(
+        self, node_id: str, start_time: Optional[str], end_time: Optional[str]
+    ) -> List[Dict[str, Any]]:
+        pass

solana_agent/interfaces/providers/graph_storage.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Any
+
+
+class GraphMemoryService(ABC):
+    """
+    Interface for a graph memory service.
+    """
+
+    @abstractmethod
+    async def add_episode(
+        self,
+        user_message: str,
+        assistant_message: str,
+        user_id: str,
+    ):
+        """
+        Add an episode to the graph memory.
+        """
+        pass
+
+    @abstractmethod
+    async def search(
+        self, query: str, user_id: str, top_k: int = 5
+    ) -> List[Dict[str, Any]]:
+        """
+        Search the graph memory for relevant episodes.
+        """
+        pass
+
+    @abstractmethod
+    async def traverse(self, node_id: str, depth: int = 1) -> List[Dict[str, Any]]:
+        """
+        Traverse the graph memory from a given node ID.
+        """
+        pass

solana_agent/interfaces/services/graph_memory.py

@@ -0,0 +1,88 @@
+import datetime
+import uuid
+from typing import Dict, Any, List
+from solana_agent.adapters.openai_adapter import OpenAIAdapter
+from solana_agent.adapters.pinecone_adapter import PineconeAdapter
+from solana_agent.adapters.mongodb_graph_adapter import MongoDBGraphAdapter
+from solana_agent.interfaces.services.graph_memory import (
+    GraphMemoryService as GraphMemoryServiceInterface,
+)
+
+
+class GraphMemoryService(GraphMemoryServiceInterface):
+    def __init__(
+        self,
+        graph_adapter: MongoDBGraphAdapter,
+        pinecone_adapter: PineconeAdapter,
+        openai_adapter: OpenAIAdapter,
+        embedding_model: str = "text-embedding-3-large",
+    ):
+        self.graph = graph_adapter
+        self.pinecone = pinecone_adapter
+        self.openai = openai_adapter
+        self.embedding_model = embedding_model
+
+    async def add_episode(
+        self,
+        user_message: str,
+        assistant_message: str,
+        user_id: str,
+    ):
+        entities = [
+            {"type": "user_message", "text": user_message, "user_id": user_id},
+            {
+                "type": "assistant_message",
+                "text": assistant_message,
+                "user_id": user_id,
+            },
+        ]
+        episode_id = str(uuid.uuid4())
+        node_ids = []
+        for entity in entities:
+            entity["episode_id"] = episode_id
+            node_id = await self.graph.add_node(entity)
+            node_ids.append(node_id)
+        edge = {
+            "source": node_ids[0],
+            "target": node_ids[1],
+            "type": "reply",
+            "episode_id": episode_id,
+            "user_id": user_id,
+            "timestamp": datetime.datetime.now(datetime.timezone.utc),
+        }
+        await self.graph.add_edge(edge)
+        # Save vectors in user-specific namespace
+        namespace = f"{user_id}_memory"
+        for node_id, entity in zip(node_ids, entities):
+            embedding = await self.openai.embed_text(
+                entity["text"], model=self.embedding_model
+            )
+            await self.pinecone.upsert(
+                [{"id": node_id, "values": embedding}],
+                namespace=namespace,
+            )
+        return episode_id
+
+    async def search(
+        self, query: str, user_id: str, top_k: int = 5
+    ) -> List[Dict[str, Any]]:
+        embedding = await self.openai.embed_text(query, model=self.embedding_model)
+        namespace = f"{user_id}_memory"
+        results = await self.pinecone.query_and_rerank(
+            vector=embedding,
+            query_text_for_rerank=query,
+            top_k=top_k,
+            namespace=namespace,
+        )
+        node_ids = [r["id"] for r in results]
+        # Only return nodes that match user_id
+        nodes = []
+        for nid in node_ids:
+            print(f"Node ID: {nid}")
+            node = await self.graph.get_node(nid)
+            if node and node.get("user_id") == user_id:
+                nodes.append(node)
+        return nodes
+
+    async def traverse(self, node_id: str, depth: int = 1) -> List[Dict[str, Any]]:
+        return await self.graph.find_neighbors(node_id, depth=depth)