solana-agent 31.0.0__py3-none-any.whl → 31.1.1__py3-none-any.whl

solana_agent/client/solana_agent.py

@@ -49,6 +49,8 @@ class SolanaAgent(SolanaAgentInterface):
         user_id: str,
         message: Union[str, bytes],
         prompt: Optional[str] = None,
+        capture_schema: Optional[Dict[str, Any]] = None,
+        capture_name: Optional[str] = None,
         output_format: Literal["text", "audio"] = "text",
         audio_voice: Literal[
             "alloy",
@@ -103,6 +105,8 @@ class SolanaAgent(SolanaAgentInterface):
             prompt=prompt,
             router=router,
             output_model=output_model,
+            capture_schema=capture_schema,
+            capture_name=capture_name,
         ):
             yield chunk
 

solana_agent/domains/agent.py

@@ -5,7 +5,7 @@ This module defines the core domain models for representing
 AI agents, human agents, and business mission/values.
 """
 
-from typing import List, Dict
+from typing import List, Dict, Optional, Any
 from pydantic import BaseModel, Field, field_validator
 
 
@@ -53,6 +53,12 @@ class AIAgent(BaseModel):
     name: str = Field(..., description="Unique agent identifier name")
     instructions: str = Field(..., description="Base instructions for the agent")
     specialization: str = Field(..., description="Agent's specialized domain")
+    capture_name: Optional[str] = Field(
+        default=None, description="Optional capture name for structured data"
+    )
+    capture_schema: Optional[Dict[str, Any]] = Field(
+        default=None, description="Optional JSON schema for structured capture"
+    )
 
     @field_validator("name", "specialization")
     @classmethod

solana_agent/factories/agent_factory.py

@@ -133,21 +133,38 @@ class SolanaAgentFactory:
                 voice=org_config.get("voice", ""),
             )
 
+        # Build capture modes from agent config if provided
+        capture_modes: Dict[str, str] = {}
+        for agent in config.get("agents", []):
+            mode = agent.get("capture_mode")
+            if mode in {"once", "multiple"} and agent.get("name"):
+                capture_modes[agent["name"]] = mode
+
         # Create repositories
         memory_provider = None
 
         if "zep" in config and "mongo" in config:
-            memory_provider = MemoryRepository(
-                mongo_adapter=db_adapter, zep_api_key=config["zep"].get("api_key")
-            )
+            mem_kwargs: Dict[str, Any] = {
+                "mongo_adapter": db_adapter,
+                "zep_api_key": config["zep"].get("api_key"),
+            }
+            if capture_modes:  # pragma: no cover
+                mem_kwargs["capture_modes"] = capture_modes
+            memory_provider = MemoryRepository(**mem_kwargs)
 
         if "mongo" in config and "zep" not in config:
-            memory_provider = MemoryRepository(mongo_adapter=db_adapter)
+            mem_kwargs = {"mongo_adapter": db_adapter}
+            if capture_modes:
+                mem_kwargs["capture_modes"] = capture_modes
+            memory_provider = MemoryRepository(**mem_kwargs)
 
         if "zep" in config and "mongo" not in config:
             if "api_key" not in config["zep"]:
                 raise ValueError("Zep API key is required.")
-            memory_provider = MemoryRepository(zep_api_key=config["zep"].get("api_key"))
+            mem_kwargs = {"zep_api_key": config["zep"].get("api_key")}
+            if capture_modes:  # pragma: no cover
+                mem_kwargs["capture_modes"] = capture_modes
+            memory_provider = MemoryRepository(**mem_kwargs)
 
         guardrail_config = config.get("guardrails", {})
         input_guardrails: List[InputGuardrail] = SolanaAgentFactory._create_guardrails(
@@ -169,9 +186,16 @@ class SolanaAgentFactory:
         )
 
         # Create routing service
+        # Optional routing model override (use small, cheap model by default in service)
+        routing_model = (
+            config.get("openai", {}).get("routing_model")
+            if isinstance(config.get("openai"), dict)
+            else None
+        )
         routing_service = RoutingService(
             llm_provider=llm_adapter,
             agent_service=agent_service,
+            model=routing_model,
         )
 
         # Debug the agent service tool registry
@@ -191,11 +215,18 @@ class SolanaAgentFactory:
             loaded_plugins = 0
 
         # Register predefined agents
-        for agent_config in config.get("agents", []):
+        for agent_config in config.get("agents", []):  # pragma: no cover
+            extra_kwargs = {}
+            if "capture_name" in agent_config:
+                extra_kwargs["capture_name"] = agent_config.get("capture_name")
+            if "capture_schema" in agent_config:
+                extra_kwargs["capture_schema"] = agent_config.get("capture_schema")
+
             agent_service.register_ai_agent(
                 name=agent_config["name"],
                 instructions=agent_config["instructions"],
                 specialization=agent_config["specialization"],
+                **extra_kwargs,
             )
 
             # Register tools for this agent

solana_agent/interfaces/providers/memory.py

@@ -36,3 +36,15 @@ class MemoryProvider(ABC):
     def count_documents(self, collection: str, query: Dict) -> int:
         """Count documents matching query."""
         pass
+
+    @abstractmethod
+    async def save_capture(
+        self,
+        user_id: str,
+        capture_name: str,
+        agent_name: Optional[str],
+        data: Dict[str, Any],
+        schema: Optional[Dict[str, Any]] = None,
+    ) -> Optional[str]:
+        """Persist a structured capture for a user and return its ID if available."""
+        pass

solana_agent/interfaces/services/query.py

@@ -38,6 +38,8 @@ class QueryService(ABC):
         router: Optional[RoutingInterface] = None,
         images: Optional[List[Union[str, bytes]]] = None,
         output_model: Optional[Type[BaseModel]] = None,
+        capture_schema: Optional[Dict[str, Any]] = None,
+        capture_name: Optional[str] = None,
     ) -> AsyncGenerator[Union[str, bytes, BaseModel], None]:
         """Process the user request and generate a response."""
         pass

solana_agent/repositories/memory.py

@@ -1,13 +1,14 @@
-import logging  # Import logging
-from copy import deepcopy
-from typing import List, Dict, Any, Optional, Tuple
+import logging
+from typing import List, Dict, Optional, Tuple, Any
 from datetime import datetime, timezone
+from copy import deepcopy
+
 from zep_cloud.client import AsyncZep as AsyncZepCloud
 from zep_cloud.types import Message
+
 from solana_agent.interfaces.providers.memory import MemoryProvider
 from solana_agent.adapters.mongodb_adapter import MongoDBAdapter
 
-# Setup logger for this module
 logger = logging.getLogger(__name__)
 
 
@@ -18,159 +19,157 @@ class MemoryRepository(MemoryProvider):
         self,
         mongo_adapter: Optional[MongoDBAdapter] = None,
         zep_api_key: Optional[str] = None,
+        capture_modes: Optional[Dict[str, str]] = None,
     ):
-        """Initialize the combined memory provider."""
+        self.capture_modes: Dict[str, str] = capture_modes or {}
+
+        # Mongo setup
         if not mongo_adapter:
             self.mongo = None
             self.collection = None
+            self.captures_collection = "captures"
         else:
-            # Initialize MongoDB
             self.mongo = mongo_adapter
             self.collection = "conversations"
-
             try:
-                # Ensure MongoDB collection and indexes
                 self.mongo.create_collection(self.collection)
                 self.mongo.create_index(self.collection, [("user_id", 1)])
                 self.mongo.create_index(self.collection, [("timestamp", 1)])
             except Exception as e:
-                logger.error(f"Error initializing MongoDB: {e}")  # Use logger.error
+                logger.error(f"Error initializing MongoDB: {e}")
+
+            try:
+                self.captures_collection = "captures"
+                self.mongo.create_collection(self.captures_collection)
+                # Basic indexes
+                self.mongo.create_index(self.captures_collection, [("user_id", 1)])
+                self.mongo.create_index(self.captures_collection, [("capture_name", 1)])
+                self.mongo.create_index(self.captures_collection, [("agent_name", 1)])
+                self.mongo.create_index(self.captures_collection, [("timestamp", 1)])
+                # Unique only when mode == 'once'
+                try:
+                    self.mongo.create_index(
+                        self.captures_collection,
+                        [("user_id", 1), ("agent_name", 1), ("capture_name", 1)],
+                        unique=True,
+                        partialFilterExpression={"mode": "once"},
+                    )
+                except Exception as e:
+                    logger.error(
+                        f"Error creating partial unique index for captures: {e}"
+                    )
+            except Exception as e:
+                logger.error(f"Error initializing MongoDB captures collection: {e}")
+                self.captures_collection = "captures"
 
-        self.zep = None
-        # Initialize Zep
-        if zep_api_key:
-            self.zep = AsyncZepCloud(api_key=zep_api_key)
+        # Zep setup
+        self.zep = AsyncZepCloud(api_key=zep_api_key) if zep_api_key else None
 
     async def store(self, user_id: str, messages: List[Dict[str, Any]]) -> None:
-        """Store messages in both Zep and MongoDB."""
-        if not user_id or user_id == "" or not isinstance(user_id, str):
+        if not user_id or not isinstance(user_id, str):
             raise ValueError("User ID cannot be None or empty")
         if not messages or not isinstance(messages, list):
             raise ValueError("Messages must be a non-empty list")
         if not all(
-            isinstance(msg, dict) and "role" in msg and "content" in msg
-            for msg in messages
+            isinstance(m, dict) and "role" in m and "content" in m for m in messages
         ):
             raise ValueError(
                 "All messages must be dictionaries with 'role' and 'content' keys"
             )
-        for msg in messages:
-            if msg["role"] not in ["user", "assistant"]:
+        for m in messages:
+            if m["role"] not in ["user", "assistant"]:
                 raise ValueError(
-                    f"Invalid role '{msg['role']}' in message. Only 'user' and 'assistant' roles are accepted."
+                    "Invalid role in message. Only 'user' and 'assistant' are accepted."
                 )
 
-        # Store in MongoDB
+        # Persist last user/assistant pair to Mongo
         if self.mongo and len(messages) >= 2:
             try:
-                # Get last user and assistant messages
                 user_msg = None
                 assistant_msg = None
-                for msg in reversed(messages):
-                    if msg.get("role") == "user" and not user_msg:
-                        user_msg = msg.get("content")
-                    elif msg.get("role") == "assistant" and not assistant_msg:
-                        assistant_msg = msg.get("content")
+                for m in reversed(messages):
+                    if m.get("role") == "user" and not user_msg:
+                        user_msg = m.get("content")
+                    elif m.get("role") == "assistant" and not assistant_msg:
+                        assistant_msg = m.get("content")
                     if user_msg and assistant_msg:
                         break
-
                 if user_msg and assistant_msg:
-                    # Store truncated messages
-                    doc = {
-                        "user_id": user_id,
-                        "user_message": user_msg,
-                        "assistant_message": assistant_msg,
-                        "timestamp": datetime.now(timezone.utc),
-                    }
-                    self.mongo.insert_one(self.collection, doc)
+                    self.mongo.insert_one(
+                        self.collection,
+                        {
+                            "user_id": user_id,
+                            "user_message": user_msg,
+                            "assistant_message": assistant_msg,
+                            "timestamp": datetime.now(timezone.utc),
+                        },
+                    )
             except Exception as e:
-                logger.error(f"MongoDB storage error: {e}")  # Use logger.error
+                logger.error(f"MongoDB storage error: {e}")
 
-        # Store in Zep
+        # Zep
         if not self.zep:
             return
 
-        # Convert messages to Zep format
-        zep_messages = []
-
-        for msg in messages:
-            if "role" in msg and "content" in msg:
-                content = self._truncate(deepcopy(msg["content"]))
-                role_type = "user" if msg["role"] == "user" else "assistant"
-                zep_msg = Message(
-                    content=content,
-                    role=role_type,
-                )
-                zep_messages.append(zep_msg)
+        zep_messages: List[Message] = []
+        for m in messages:
+            content = (
+                self._truncate(deepcopy(m.get("content"))) if "content" in m else None
+            )
+            if content is None:  # pragma: no cover
+                continue
+            role_type = "user" if m.get("role") == "user" else "assistant"
+            zep_messages.append(Message(content=content, role=role_type))
 
-        # Add messages to Zep memory
         if zep_messages:
             try:
                 await self.zep.thread.add_messages(
-                    thread_id=user_id,
-                    messages=zep_messages,
+                    thread_id=user_id, messages=zep_messages
                 )
             except Exception:
                 try:
                     try:
                         await self.zep.user.add(user_id=user_id)
                     except Exception as e:
-                        logger.error(
-                            f"Zep user addition error: {e}"
-                        )  # Use logger.error
-
+                        logger.error(f"Zep user addition error: {e}")
                     try:
-                        await self.zep.thread.create(
-                            thread_id=user_id,
-                            user_id=user_id,
-                        )
+                        await self.zep.thread.create(thread_id=user_id, user_id=user_id)
                     except Exception as e:
-                        logger.error(
-                            f"Zep thread creation error: {e}"
-                        )  # Use logger.error
+                        logger.error(f"Zep thread creation error: {e}")
                     await self.zep.thread.add_messages(
-                        thread_id=user_id,
-                        messages=zep_messages,
+                        thread_id=user_id, messages=zep_messages
                     )
                 except Exception as e:
-                    logger.error(f"Zep memory addition error: {e}")  # Use logger.error
-                    return
+                    logger.error(f"Zep memory addition error: {e}")
 
     async def retrieve(self, user_id: str) -> str:
-        """Retrieve memory context from Zep."""
         try:
             memories = ""
             if self.zep:
                 memory = await self.zep.thread.get_user_context(thread_id=user_id)
                 if memory and memory.context:
                     memories = memory.context
-
             return memories
-
         except Exception as e:
-            logger.error(f"Error retrieving memories: {e}")  # Use logger.error
+            logger.error(f"Error retrieving memories: {e}")
             return ""
 
     async def delete(self, user_id: str) -> None:
-        """Delete memory from both systems."""
         if self.mongo:
             try:
                 self.mongo.delete_all(self.collection, {"user_id": user_id})
             except Exception as e:
-                logger.error(f"MongoDB deletion error: {e}")  # Use logger.error
-
+                logger.error(f"MongoDB deletion error: {e}")
         if not self.zep:
             return
-
         try:
             await self.zep.thread.delete(thread_id=user_id)
         except Exception as e:
-            logger.error(f"Zep memory deletion error: {e}")  # Use logger.error
-
+            logger.error(f"Zep memory deletion error: {e}")
         try:
             await self.zep.user.delete(user_id=user_id)
         except Exception as e:
-            logger.error(f"Zep user deletion error: {e}")  # Use logger.error
+            logger.error(f"Zep user deletion error: {e}")
 
     def find(
         self,
@@ -180,37 +179,98 @@ class MemoryRepository(MemoryProvider):
         limit: int = 0,
         skip: int = 0,
     ) -> List[Dict]:  # pragma: no cover
-        """Find documents in MongoDB."""
         if not self.mongo:
             return []
-
         try:
             return self.mongo.find(collection, query, sort=sort, limit=limit, skip=skip)
         except Exception as e:
-            logger.error(f"MongoDB find error: {e}")  # Use logger.error
+            logger.error(f"MongoDB find error: {e}")
             return []
 
     def count_documents(self, collection: str, query: Dict) -> int:
-        """Count documents in MongoDB."""
         if not self.mongo:
             return 0
         return self.mongo.count_documents(collection, query)
 
     def _truncate(self, text: str, limit: int = 2500) -> str:
-        """Truncate text to be within limits."""
         if text is None:
             raise AttributeError("Cannot truncate None text")
-
         if not text:
             return ""
-
         if len(text) <= limit:
             return text
-
-        # Try to truncate at last period before limit
         last_period = text.rfind(".", 0, limit)
         if last_period > 0:
             return text[: last_period + 1]
-
-        # If no period found, truncate at limit and add ellipsis
         return text[: limit - 3] + "..."
+
+    async def save_capture(
+        self,
+        user_id: str,
+        capture_name: str,
+        agent_name: Optional[str],
+        data: Dict[str, Any],
+        schema: Optional[Dict[str, Any]] = None,
+    ) -> Optional[str]:
+        if not self.mongo:  # pragma: no cover
+            logger.warning("MongoDB not configured; cannot save capture.")
+            return None
+        if not user_id or not isinstance(user_id, str):
+            raise ValueError("user_id must be a non-empty string")
+        if not capture_name or not isinstance(capture_name, str):
+            raise ValueError("capture_name must be a non-empty string")
+        if not isinstance(data, dict):
+            raise ValueError("data must be a dictionary")
+
+        try:
+            mode = self.capture_modes.get(agent_name, "once") if agent_name else "once"
+            now = datetime.now(timezone.utc)
+            if mode == "multiple":
+                doc = {
+                    "user_id": user_id,
+                    "agent_name": agent_name,
+                    "capture_name": capture_name,
+                    "data": data or {},
+                    "schema": schema or {},
+                    "mode": "multiple",
+                    "timestamp": now,
+                    "created_at": now,
+                }
+                return self.mongo.insert_one(self.captures_collection, doc)
+            else:
+                key = {
+                    "user_id": user_id,
+                    "agent_name": agent_name,
+                    "capture_name": capture_name,
+                }
+                existing = self.mongo.find_one(self.captures_collection, key)
+                merged_data: Dict[str, Any] = {}
+                if existing and isinstance(existing.get("data"), dict):
+                    merged_data.update(existing.get("data", {}))
+                merged_data.update(data or {})
+                update_doc = {
+                    "$set": {
+                        "user_id": user_id,
+                        "agent_name": agent_name,
+                        "capture_name": capture_name,
+                        "data": merged_data,
+                        "schema": (
+                            schema
+                            if schema is not None
+                            else existing.get("schema")
+                            if existing
+                            else {}
+                        ),
+                        "mode": "once",
+                        "timestamp": now,
+                    },
+                    "$setOnInsert": {"created_at": now},
+                }
+                self.mongo.update_one(
+                    self.captures_collection, key, update_doc, upsert=True
+                )
+                doc = self.mongo.find_one(self.captures_collection, key)
+                return str(doc.get("_id")) if doc and doc.get("_id") else None
+        except Exception as e:  # pragma: no cover
+            logger.error(f"MongoDB save_capture error: {e}")
+            return None

solana_agent/services/agent.py

@@ -73,6 +73,8 @@ class AgentService(AgentServiceInterface):
         name: str,
         instructions: str,
         specialization: str,
+        capture_name: Optional[str] = None,
+        capture_schema: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Register an AI agent with its specialization.
 
@@ -85,6 +87,8 @@ class AgentService(AgentServiceInterface):
             name=name,
             instructions=instructions,
             specialization=specialization,
+            capture_name=capture_name,
+            capture_schema=capture_schema,
         )
         self.agents.append(agent)
         logger.info(f"Registered AI agent: {name}")
@@ -98,7 +102,6 @@ class AgentService(AgentServiceInterface):
         Returns:
             System prompt
         """
-
         # Get agent by name
         agent = next((a for a in self.agents if a.name == agent_name), None)
 
@@ -130,8 +133,37 @@ class AgentService(AgentServiceInterface):
                 )
                 system_prompt += f"\n\nBUSINESS GOALS:\n{goals_text}"
 
+        # Add capture guidance if this agent has a capture schema
+        if getattr(agent, "capture_schema", None) and getattr(
+            agent, "capture_name", None
+        ):  # pragma: no cover
+            system_prompt += (
+                "\n\nSTRUCTURED DATA CAPTURE:\n"
+                f"You must collect the following fields for the form '{agent.capture_name}'. "
+                "Ask concise follow-up questions to fill any missing required fields one at a time. "
+                "Confirm values when ambiguous, and summarize the captured data before finalizing.\n\n"
+                "JSON Schema (authoritative definition of the fields):\n"
+                f"{agent.capture_schema}\n\n"
+                "Rules:\n"
+                "- Never invent values—ask the user.\n"
+                "- Validate types (emails look like emails, numbers are numbers, booleans are yes/no).\n"
+                "- If the user declines to provide a required value, note it clearly.\n"
+                "- When all required fields are provided, acknowledge completion.\n"
+            )
+
         return system_prompt
 
+    def get_agent_capture(
+        self, agent_name: str
+    ) -> Optional[Dict[str, Any]]:  # pragma: no cover
+        """Return capture metadata for the agent, if any."""
+        agent = next((a for a in self.agents if a.name == agent_name), None)
+        if not agent:
+            return None
+        if agent.capture_name and agent.capture_schema:
+            return {"name": agent.capture_name, "schema": agent.capture_schema}
+        return None
+
     def get_all_ai_agents(self) -> Dict[str, AIAgent]:
         """Get all registered AI agents.