bump version to 0.0.5 and implement memory storage classes

Author: Rakesh-Sahoo77

build/lib/semantio/agent.py

@@ -16,6 +16,7 @@
 from pathlib import Path
 import importlib
 import os
+from .memory import Memory
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -48,6 +49,13 @@ class Agent(BaseModel):
     semantic_model: Optional[Any] = Field(None, description="SentenceTransformer model for semantic matching.")
     team: Optional[List['Agent']] = Field(None, description="List of assistants in the team.")
     auto_tool: bool = Field(False, description="Whether to automatically detect and call tools.")
+    memory: Memory = Field(default_factory=Memory)
+    memory_config: Dict = Field(
+        default_factory=lambda: {
+            "max_context_length": 4000,
+            "summarization_threshold": 3000
+        }
+    )
 
     # Allow arbitrary types
     model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -56,6 +64,11 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
         # Initialize the model and tools here if needed
         self._initialize_model()
+        # Initialize memory with config
+        self.memory = Memory(
+            max_context_length=self.memory_config.get("max_context_length", 4000),
+            summarization_threshold=self.memory_config.get("summarization_threshold", 3000)
+        )
         # Initialize tools as an empty list if not provided
         if self.tools is None:
             self.tools = []
@@ -218,20 +231,31 @@ def print_response(
         markdown: bool = False,
         team: Optional[List['Agent']] = None,
         **kwargs,
-    ) -> Union[str, Dict]:  # Add return type hint
+    ) -> Union[str, Dict]:
         """Print the agent's response to the console and return it."""
+    
+        # Store user message if provided
+        if message and isinstance(message, str):
+            self.memory.add_message(role="user", content=message)
 
         if stream:
             # Handle streaming response
             response = ""
             for chunk in self._stream_response(message, markdown=markdown, **kwargs):
-                print(chunk)
+                print(chunk, end="", flush=True)
                 response += chunk
+            # Store agent response
+            if response:
+                self.memory.add_message(role="assistant", content=response)
+            print()  # New line after streaming
             return response
         else:
             # Generate and return the response
             response = self._generate_response(message, markdown=markdown, team=team, **kwargs)
             print(response)  # Print the response to the console
+            # Store agent response
+            if response:
+                self.memory.add_message(role="assistant", content=response)
             return response
 
 
@@ -294,12 +318,10 @@ def _generate_response(self, message: str, markdown: bool = False, team: Optiona
         # Use the specified team if provided
         if team is not None:
             return self._generate_team_response(message, team, markdown=markdown, **kwargs)
-
         # Initialize tool_outputs as an empty dictionary
         tool_outputs = {}
         responses = []
         tool_calls = []
-
         # Use the LLM to analyze the query and dynamically select tools when auto_tool is enabled
         if self.auto_tool:
             tool_calls = self._analyze_query_and_select_tools(message)
@@ -347,13 +369,17 @@ def _generate_response(self, message: str, markdown: bool = False, team: Optiona
             try:
                 # Prepare the context for the LLM
                 context = {
+                    "conversation_history": self.memory.get_context(self.llm_instance),
                     "tool_outputs": tool_outputs,
                     "rag_context": self.rag.retrieve(message) if self.rag else None,
-                    "knowledge_base_context": self._find_all_relevant_keys(message, self._flatten_data(self.knowledge_base)) if self.knowledge_base else None,
+                    "knowledge_base": self._get_knowledge_context(message) if self.knowledge_base else None,
                 }
-
+                # 3. Build a memory-aware prompt.
+                prompt = self._build_memory_prompt(message, context)
+                # To (convert MemoryEntry objects to dicts and remove metadata):
+                memory_entries = [{"role": e.role, "content": e.content} for e in self.memory.storage.retrieve()]
                 # Generate a response using the LLM
-                llm_response = self.llm_instance.generate(prompt=message, context=context, **kwargs)
+                llm_response = self.llm_instance.generate(prompt=prompt, context=context, memory=memory_entries, **kwargs)
                 responses.append(f"**Analysis:**\n\n{llm_response}")
             except Exception as e:
                 logger.error(f"Failed to generate LLM response: {e}")
@@ -363,25 +389,30 @@ def _generate_response(self, message: str, markdown: bool = False, team: Optiona
             # Retrieve relevant context using RAG
             rag_context = self.rag.retrieve(message) if self.rag else None
             # Retrieve relevant context from the knowledge base (API result)
-            knowledge_base_context = None
-            if self.knowledge_base:
-                # Flatten the knowledge base
-                flattened_data = self._flatten_data(self.knowledge_base)
-                # Find all relevant key-value pairs in the knowledge base
-                relevant_values = self._find_all_relevant_keys(message, flattened_data)
-                if relevant_values:
-                    knowledge_base_context = ", ".join(relevant_values)
+            # knowledge_base_context = None
+            # if self.knowledge_base:
+            #     # Flatten the knowledge base
+            #     flattened_data = self._flatten_data(self.knowledge_base)
+            #     # Find all relevant key-value pairs in the knowledge base
+            #     relevant_values = self._find_all_relevant_keys(message, flattened_data)
+            #     if relevant_values:
+            #         knowledge_base_context = ", ".join(relevant_values)
 
             # Combine both contexts (RAG and knowledge base)
             context = {
+                "conversation_history": self.memory.get_context(self.llm_instance),
                 "rag_context": rag_context,
-                "knowledge_base_context": knowledge_base_context,
+                "knowledge_base": self._get_knowledge_context(message),
             }
             # Prepare the prompt with instructions, description, and context
-            prompt = self._build_prompt(message, context)
+            # 3. Build a memory-aware prompt.
+            prompt = self._build_memory_prompt(message, context)
+            # To (convert MemoryEntry objects to dicts and remove metadata):
+            memory_entries = [{"role": e.role, "content": e.content} for e in self.memory.storage.retrieve()]
 
             # Generate the response using the LLM
-            response = self.llm_instance.generate(prompt=prompt, context=context, **kwargs)
+            response = self.llm_instance.generate(prompt=prompt, context=context, memory=memory_entries, **kwargs)
+
 
             # Format the response based on the json_output flag
             if self.json_output:
@@ -394,9 +425,37 @@ def _generate_response(self, message: str, markdown: bool = False, team: Optiona
             if markdown:
                 return f"**Response:**\n\n{response}"
             return response
-        # Combine all responses into a single string
         return "\n\n".join(responses)
 
+    # Modified prompt construction with memory integration
+    def _build_memory_prompt(self, user_input: str, context: dict) -> str:
+        """Enhanced prompt builder with memory context."""
+        prompt_parts = []
+        
+        if self.description:
+            prompt_parts.append(f"# ROLE\n{self.description}")
+            
+        if self.instructions:
+            prompt_parts.append(f"# INSTRUCTIONS\n" + "\n".join(f"- {i}" for i in self.instructions))
+            
+        if context['conversation_history']:
+            prompt_parts.append(f"# CONVERSATION HISTORY\n{context['conversation_history']}")
+            
+        if context['knowledge_base']:
+            prompt_parts.append(f"# KNOWLEDGE BASE\n{context['knowledge_base']}")
+            
+        prompt_parts.append(f"# USER INPUT\n{user_input}")
+        
+        return "\n\n".join(prompt_parts)
+        
+    def _get_knowledge_context(self, message: str) -> str:
+        """Retrieve and format knowledge base context."""
+        if not self.knowledge_base:
+            return ""
+        
+        flattened = self._flatten_data(self.knowledge_base)
+        relevant = self._find_all_relevant_keys(message, flattened)
+        return "\n".join(f"- {item}" for item in relevant) if relevant else ""
     def _generate_team_response(self, message: str, team: List['Agent'], markdown: bool = False, **kwargs) -> str:
         """Generate a response using a team of assistants."""
         responses = []
@@ -543,17 +602,21 @@ def cli_app(
         """Run the agent in a CLI app."""
         from rich.prompt import Prompt
 
+        # Print initial message if provided
         if message:
             self.print_response(message=message, **kwargs)
 
         _exit_on = exit_on or ["exit", "quit", "bye"]
         while True:
-            message = Prompt.ask(f"[bold] {self.emoji} {self.user_name} [/bold]")
-            if message in _exit_on:
+            try:
+                message = Prompt.ask(f"[bold] {self.emoji} {self.user_name} [/bold]")
+                if message in _exit_on:
+                    break
+                self.print_response(message=message, **kwargs)
+            except KeyboardInterrupt:
+                print("\n\nSession ended. Goodbye!")
                 break
 
-            self.print_response(message=message, **kwargs)
-
     def _generate_api(self):
         """Generate an API for the agent if api=True."""
         from .api.api_generator import APIGenerator

build/lib/semantio/memory.py

@@ -1,11 +1,54 @@
-from typing import List, Dict
-
+from .models import MemoryEntry
+from .storage import BaseMemoryStorage, InMemoryStorage, FileStorage
+from typing import List, Dict, Optional
+from .llm.base_llm import BaseLLM
 class Memory:
-    def __init__(self):
-        self.history = []
+    def __init__(
+        self,
+        storage: BaseMemoryStorage = InMemoryStorage(),
+        max_context_length: int = 4000,
+        summarization_threshold: int = 3000
+    ):
+        self.storage = storage
+        self.max_context_length = max_context_length
+        self.summarization_threshold = summarization_threshold
+        self._current_context = ""
+
+    def add_message(self, role: str, content: str, metadata: Optional[Dict] = None):
+        entry = MemoryEntry(
+            role=role,
+            content=content,
+            metadata=metadata or {}
+        )
+        self.storage.store(entry)
+        self._manage_context()
+
+    def get_context(self, llm: Optional[BaseLLM] = None) -> str:
+        if len(self._current_context) < self.summarization_threshold:
+            return self._current_context
+        
+        # Automatic summarization when context grows too large
+        if llm:
+            return self.summarize(llm)
+        return self._current_context[:self.max_context_length]
+    def _manage_context(self):
+        # Include roles in the conversation history
+        full_history = "\n".join([f"{e.role}: {e.content}" for e in self.storage.retrieve()])
+        if len(full_history) > self.max_context_length:
+            self._current_context = full_history[-self.max_context_length:]
+        else:
+            self._current_context = full_history
 
-    def add_message(self, role: str, content: str):
-        self.history.append({"role": role, "content": content})
+    def summarize(self, llm: BaseLLM) -> str:
+        # Include roles in the history for summarization
+        history = "\n".join([f"{e.role}: {e.content}" for e in self.storage.retrieve()])
+        prompt = f"""
+        Summarize this conversation history maintaining key details and references:
+        {history[-self.summarization_threshold:]}
+        """
+        self._current_context = llm.generate(prompt)
+        return self._current_context
 
-    def get_history(self) -> List[Dict]:
-        return self.history
+    def clear(self):
+        self.storage = InMemoryStorage()
+        self._current_context = ""

build/lib/semantio/models.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, Field
+from datetime import datetime
+from typing import Dict
+
+class MemoryEntry(BaseModel):
+    role: str  # "user" or "assistant"
+    content: str
+    timestamp: datetime = Field(default_factory=datetime.now)
+    metadata: Dict = Field(default_factory=dict)

build/lib/semantio/storage/__init__.py

@@ -0,0 +1,5 @@
+from .base_storage import BaseMemoryStorage
+from .in_memory_storage import InMemoryStorage
+from .local_storage import FileStorage
+
+__all__ = ['BaseMemoryStorage', 'InMemoryStorage', 'FileStorage']

build/lib/semantio/storage/base_storage.py

@@ -0,0 +1,12 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+from ..models import MemoryEntry
+
+class BaseMemoryStorage(ABC):
+    @abstractmethod
+    def store(self, entry: MemoryEntry):
+        pass
+    
+    @abstractmethod
+    def retrieve(self, query: Optional[str] = None, limit: int = 20) -> List[MemoryEntry]:
+        pass

build/lib/semantio/storage/in_memory_storage.py

@@ -0,0 +1,14 @@
+# hashai/storage/in_memory_storage.py
+from typing import List, Optional
+from ..models import MemoryEntry
+from .base_storage import BaseMemoryStorage
+
+class InMemoryStorage(BaseMemoryStorage):
+    def __init__(self):
+        self.history: List[MemoryEntry] = []
+    
+    def store(self, entry: MemoryEntry):
+        self.history.append(entry)
+    
+    def retrieve(self, query: Optional[str] = None, limit: int = 10) -> List[MemoryEntry]:
+        return self.history[-limit:]

build/lib/semantio/storage/local_storage.py

@@ -0,0 +1,29 @@
+import json
+from typing import List, Optional
+from ..models import MemoryEntry
+from .base_storage import BaseMemoryStorage
+
+class FileStorage(BaseMemoryStorage):
+    def __init__(self, file_path: str = "memory.json"):
+        self.file_path = file_path
+        self.history = self._load_from_file()
+
+    def _load_from_file(self) -> List[MemoryEntry]:
+        try:
+            with open(self.file_path, "r") as f:
+                data = json.load(f)
+                return [MemoryEntry(**entry) for entry in data]
+        except (FileNotFoundError, json.JSONDecodeError):
+            return []
+
+    def _save_to_file(self):
+        with open(self.file_path, "w") as f:
+            data = [entry.dict() for entry in self.history]
+            json.dump(data, f, default=str)
+
+    def store(self, entry: MemoryEntry):
+        self.history.append(entry)
+        self._save_to_file()
+
+    def retrieve(self, query: Optional[str] = None, limit: int = 20) -> List[MemoryEntry]:
+        return self.history[-limit:]

dist/semantio-0.0.5-py3-none-any.whl

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: semantio
-Version: 0.0.4
+Version: 0.0.5
 Summary: A powerful SDK for building AI agents
 Home-page: https://github.com/Syenah/semantio
 Author: Rakesh