litellm-adk 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

litellm_adk/__init__.py

@@ -1,10 +1,12 @@
 from .agents import LiteLLMAgent
+from .session import Session
 from .tools import tool, tool_registry
 from .config.settings import settings
 from .memory import BaseMemory, InMemoryMemory, FileMemory, MongoDBMemory
 
 __all__ = [
     "LiteLLMAgent", 
+    "Session",
     "tool", 
     "tool_registry", 
     "settings", 

litellm_adk/core/agent.py

@@ -4,9 +4,11 @@ from typing import List, Dict, Any, Optional, Union, Generator, AsyncGenerator
 from .base import BaseAgent
 from ..observability.logger import adk_logger
 from ..config.settings import settings
+from ..session import Session
 from ..tools.registry import tool_registry
 from ..memory.base import BaseMemory
 from ..memory.in_memory import InMemoryMemory
+from .context import ContextManager
 import uuid
 
 # Global LiteLLM configuration for resilience
@@ -25,7 +27,7 @@ class LiteLLMAgent(BaseAgent):
         system_prompt: str = "You are a helpful assistant.",
         tools: Optional[List[Dict[str, Any]]] = None,
         memory: Optional[BaseMemory] = None,
-        session_id: Optional[str] = None,
+        max_context_tokens: Optional[int] = None,
         **kwargs
     ):
         self.model = model or settings.model
@@ -57,71 +59,103 @@ class LiteLLMAgent(BaseAgent):
 
         self.extra_kwargs = kwargs
         
-        # Ensure model-specific parameters        # Default parallel_tool_calls if not explicitly provided
         if "parallel_tool_calls" not in self.extra_kwargs:
-            self.extra_kwargs["parallel_tool_calls"] = False
+            self.extra_kwargs["parallel_tool_calls"] = True
             
         self.sequential_tool_execution = kwargs.get("sequential_tool_execution", settings.sequential_execution)
         
         # Memory Persistence
         self.memory = memory or InMemoryMemory()
-        self.session_id = session_id or str(uuid.uuid4())
-        self.history = self.memory.get_messages(self.session_id)
+        self.max_context_tokens = max_context_tokens
         
-        if not self.history:
-            self.history = [{"role": "system", "content": self.system_prompt}]
-            self.memory.add_message(self.session_id, self.history[0])
-            
-        adk_logger.debug(f"Initialized LiteLLMAgent with session_id={self.session_id}, model={self.model}")
+        adk_logger.debug(f"Initialized LiteLLMAgent as a service for model={self.model}")
 
-    def _prepare_messages(self, prompt: str) -> List[Dict[str, str]]:
-        # Refresh from memory in case it was modified elsewhere
-        self.history = self.memory.get_messages(self.session_id)
+    def save_session(self, session: Union[str, Session]):
+        """Persist session metadata and state to memory."""
+        actual_id = session.id if isinstance(session, Session) else session
         
-        messages = self.history.copy()
+        # If it's a Session object, we dump the full metadata
+        if isinstance(session, Session):
+            self.memory.save_session_metadata(actual_id, session.model_dump())
+        # If it's just an ID, there's nothing to dump from the service layer
+
+    def _prepare_messages(self, prompt: str, actual_session_id: str) -> List[Dict[str, str]]:
+        # 2. Fetch/Initialize History from Memory
+        history = self.memory.get_messages(actual_session_id)
+        if not history:
+            history = [{"role": "system", "content": self.system_prompt}]
+            # Don't persist system prompt until first real turn to keep DB clean
+            
+        messages = history.copy()
         user_msg = {"role": "user", "content": prompt}
         messages.append(user_msg)
         
-        # Persist the user message immediately
-        self.memory.add_message(self.session_id, user_msg)
-        self.history.append(user_msg)
+        # 3. Persist turn start
+        # Ensure messages are sanitized and tokenized before first persistence
+        current_user_msg = self._sanitize_message(user_msg)
+        current_user_msg["token_count"] = ContextManager.count_tokens([current_user_msg], self.model)
         
+        if not history:
+            system_msg = self._sanitize_message(messages[0])
+            system_msg["token_count"] = ContextManager.count_tokens([system_msg], self.model)
+            self.memory.add_messages(actual_session_id, [system_msg, current_user_msg])
+        else:
+            self.memory.add_message(actual_session_id, current_user_msg)
+        
+        # 4. Context Management (Truncation)
+        if self.max_context_tokens:
+            messages = ContextManager.truncate_history(
+                messages, 
+                self.model, 
+                self.max_context_tokens
+            )
+            
         return messages
 
-    def _update_history(self, final_messages: List[Dict[str, Any]]):
-        """Sync internal history and memory with the final message state."""
-        # Find which messages were added since we prepared (the user message was already added)
-        # We assume messages order is preserved
-        start_idx = len(self.history)
-        new_messages = [self._sanitize_message(m) for m in final_messages[start_idx:]]
-        
+    def _update_history(self, new_messages: List[Dict[str, Any]], actual_session_id: str):
+        """Persist new messages to memory with token counts."""
         if new_messages:
-            self.memory.add_messages(self.session_id, new_messages)
-            self.history.extend(new_messages)
+            sanitized = []
+            for m in new_messages:
+                s = self._sanitize_message(m)
+                # Compute token count if not already present (optimization for future turns)
+                if "token_count" not in s:
+                    s["token_count"] = ContextManager.count_tokens([s], self.model)
+                sanitized.append(s)
+                
+            self.memory.add_messages(actual_session_id, sanitized)
 
     def _sanitize_message(self, message: Any) -> Dict[str, Any]:
-        """Convert LiteLLM message objects to plain dictionaries for serialization."""
-        if isinstance(message, dict):
-            # Still need to sanitize tool_calls inside if they are objects
-            if "tool_calls" in message and message["tool_calls"]:
-                message["tool_calls"] = [self._sanitize_tool_call(tc) for tc in message["tool_calls"]]
-            return message
+        """
+        Convert LiteLLM message objects to strictly compliant dictionaries.
+        Ensures compatibility with strict providers like OCI.
+        """
+        # If it's already a dict, extract only what we need to avoid 'extra key' errors
+        role = getattr(message, "role", "assistant") if not isinstance(message, dict) else message.get("role", "assistant")
+        content = getattr(message, "content", "") if not isinstance(message, dict) else message.get("content", "")
         
-        # Manually extract common fields to ensure clean JSON
+        # OCI/OpenAI standard: content cannot be None for assistant/user/system
+        if content is None:
+            content = ""
+            
         msg_dict = {
-            "role": getattr(message, "role", "assistant"),
-            "content": getattr(message, "content", None)
+            "role": role,
+            "content": content
         }
         
-        if hasattr(message, "name") and message.name:
-            msg_dict["name"] = message.name
-            
-        if hasattr(message, "tool_calls") and message.tool_calls:
-            msg_dict["tool_calls"] = [self._sanitize_tool_call(tc) for tc in message.tool_calls]
-            
-        if hasattr(message, "tool_call_id") and message.tool_call_id:
-            msg_dict["tool_call_id"] = message.tool_call_id
+        # Handle Tool Calls (Assistant Message)
+        tool_calls = getattr(message, "tool_calls", None) if not isinstance(message, dict) else message.get("tool_calls")
+        if tool_calls:
+            msg_dict["tool_calls"] = [self._sanitize_tool_call(tc) for tc in tool_calls]
             
+        # Handle Tool Result (Tool Role)
+        if role == "tool":
+            msg_dict["tool_call_id"] = getattr(message, "tool_call_id", None) if not isinstance(message, dict) else message.get("tool_call_id")
+            # Name is optional but good practice
+            name = getattr(message, "name", None) if not isinstance(message, dict) else message.get("name")
+            if name:
+                msg_dict["name"] = name
+                
         return msg_dict
 
     def _sanitize_tool_call(self, tc: Any) -> Dict[str, Any]:
@@ -149,9 +183,28 @@ class LiteLLMAgent(BaseAgent):
         """Determines if we should process tool calls one by one."""
         return self.sequential_tool_execution
 
-    async def _aexecute_tool(self, tool_call) -> Any:
-        # Same as _execute_tool but for async if needed in future
-        return self._execute_tool(tool_call)
+    async def _aexecute_tool(self, tool_call) -> Dict[str, Any]:
+        """Helper to execute a tool call asynchronously and return formatted result."""
+        function_name = self._get_tc_val(tool_call, "function", "name")
+        raw_args = self._get_tc_val(tool_call, "function", "arguments") or "{}"
+        
+        try:
+            if isinstance(raw_args, dict):
+                arguments = raw_args
+            else:
+                arguments = json.loads(raw_args)
+        except json.JSONDecodeError:
+            adk_logger.warning(f"Failed to parse tool arguments for {function_name}: {raw_args}")
+            arguments = {}
+            
+        result = await tool_registry.aexecute(function_name, **arguments)
+        
+        return {
+            "role": "tool",
+            "tool_call_id": self._get_tc_val(tool_call, "id"),
+            "name": function_name,
+            "content": str(result)
+        }
 
     def _get_tc_val(self, tool_call, attr, subattr=None):
         """Helper to get value from either object or dict tool call."""
@@ -193,12 +246,14 @@ class LiteLLMAgent(BaseAgent):
         
         return tool_registry.execute(function_name, **arguments)
 
-    def invoke(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> str:
+    def invoke(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, session_id: Optional[Union[str, Session]] = None, **kwargs) -> str:
         """
         Execute a synchronous completion with automatic tool calling.
         """
-        messages = self._prepare_messages(prompt)
+        actual_session_id = session_id.id if isinstance(session_id, Session) else (session_id or str(uuid.uuid4()))
+        messages = self._prepare_messages(prompt, actual_session_id=actual_session_id)
         tools = tools or self.tools
+        new_turns = [] # Track only what's new in this specific call
         
         adk_logger.info(f"Invoking completion for model: {self.model}")
         
@@ -216,39 +271,45 @@ class LiteLLMAgent(BaseAgent):
             
             # Check if the model wants to call tools
             if hasattr(message, "tool_calls") and message.tool_calls:
-                # If sequential is enabled, we only process the FIRST tool call
                 tool_calls_to_process = [message.tool_calls[0]] if self._should_handle_sequentially() else message.tool_calls
                 
-                # We update the original message to only include the calls we are handling 
-                # (to keep history clean for strict models)
                 if self._should_handle_sequentially():
                     message.tool_calls = tool_calls_to_process
 
                 sanitized_msg = self._sanitize_message(message)
                 messages.append(sanitized_msg)
+                new_turns.append(sanitized_msg)
                 
                 for tool_call in tool_calls_to_process:
                     result = self._execute_tool(tool_call)
                     
-                    messages.append({
+                    tool_response = {
                         "role": "tool",
                         "tool_call_id": tool_call.id,
                         "name": tool_call.function.name,
                         "content": str(result)
-                    })
+                    }
+                    messages.append(tool_response)
+                    new_turns.append(tool_response)
                 
                 continue
             
-            messages.append(self._sanitize_message(message))
-            self._update_history(messages)
-            return message.content
+            final_msg = self._sanitize_message(message)
+            messages.append(final_msg)
+            new_turns.append(final_msg)
+            
+            # Persist only the new assistant/tool turns
+            self._update_history(new_turns, actual_session_id=actual_session_id)
+            return final_msg.get("content")
 
-    async def ainvoke(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> str:
+    async def ainvoke(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, session_id: Optional[Union[str, Session]] = None, **kwargs) -> str:
         """
         Execute an asynchronous completion with automatic tool calling.
         """
-        messages = self._prepare_messages(prompt)
+        actual_session_id = session_id.id if isinstance(session_id, Session) else (session_id or str(uuid.uuid4()))
+        messages = self._prepare_messages(prompt, actual_session_id=actual_session_id)
         tools = tools or self.tools
+        new_turns = []
         
         adk_logger.info(f"Invoking async completion for model: {self.model}")
         
@@ -272,28 +333,39 @@ class LiteLLMAgent(BaseAgent):
 
                 sanitized_msg = self._sanitize_message(message)
                 messages.append(sanitized_msg)
+                new_turns.append(sanitized_msg)
                 
-                for tool_call in tool_calls_to_process:
-                    result = self._execute_tool(tool_call)
-                    messages.append({
-                        "role": "tool",
-                        "tool_call_id": tool_call.id,
-                        "name": tool_call.function.name,
-                        "content": str(result)
-                    })
+                if self._should_handle_sequentially():
+                    for tool_call in tool_calls_to_process:
+                        result = await self._aexecute_tool(tool_call)
+                        messages.append(result)
+                        new_turns.append(result)
+                else:
+                    # Parallel Execution
+                    import asyncio
+                    results = await asyncio.gather(*[self._aexecute_tool(tc) for tc in tool_calls_to_process])
+                    for res in results:
+                        messages.append(res)
+                        new_turns.append(res)
                 continue
             
-            messages.append(self._sanitize_message(message))
-            self._update_history(messages)
-            return message.content
+            final_msg = self._sanitize_message(message)
+            messages.append(final_msg)
+            new_turns.append(final_msg)
+            
+            self._update_history(new_turns, actual_session_id=actual_session_id)
+            return final_msg.get("content")
 
-    def stream(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> Generator[str, None, None]:
+    def stream(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, session_id: Optional[Union[str, Session]] = None, **kwargs) -> Generator[str, None, None]:
         """
         Execute a streaming completion with automatic tool calling.
         """
-        messages = self._prepare_messages(prompt)
+        actual_session_id = session_id.id if isinstance(session_id, Session) else (session_id or str(uuid.uuid4()))
+        messages = self._prepare_messages(prompt, actual_session_id=actual_session_id)
         tools = tools or self.tools
         
+        new_turns = []
+        
         while True:
             response = litellm.completion(
                 model=self.model,
@@ -356,7 +428,7 @@ class LiteLLMAgent(BaseAgent):
                                     if last_tc.function.arguments is None:
                                         last_tc.function.arguments = ""
                                     last_tc.function.arguments += tc_delta.function.arguments
-
+  
             # Build final flattened tool calls list (as dicts for history)
             tool_calls = []
             for idx in sorted(tool_calls_by_index.keys()):
@@ -370,39 +442,46 @@ class LiteLLMAgent(BaseAgent):
                                 "arguments": tc_obj.function.arguments
                             }
                         })
-
+  
             if tool_calls:
                 # If sequential, only keep the first tool call
                 if self._should_handle_sequentially():
                     tool_calls = [tool_calls[0]]
-
+  
                 # Add the assistant's composite tool call message to history
-                assistant_msg = {"role": "assistant", "tool_calls": tool_calls, "content": full_content or None}
+                assistant_msg = self._sanitize_message({"role": "assistant", "tool_calls": tool_calls, "content": full_content})
                 messages.append(assistant_msg)
+                new_turns.append(assistant_msg)
                 
                 for tool_call in tool_calls:
                     result = self._execute_tool(tool_call)
-                    messages.append({
+                    tool_resp = {
                         "role": "tool",
                         "tool_call_id": tool_call["id"],
                         "name": tool_call["function"]["name"],
                         "content": str(result)
-                    })
+                    }
+                    messages.append(tool_resp)
+                    new_turns.append(tool_resp)
                 
                 # Loop back to continue the conversation with tool results
                 continue
             
             # No tool calls, store final content and finish
-            messages.append({"role": "assistant", "content": full_content})
-            self._update_history(messages)
+            final_msg = self._sanitize_message({"role": "assistant", "content": full_content})
+            messages.append(final_msg)
+            new_turns.append(final_msg)
+            self._update_history(new_turns, actual_session_id=actual_session_id)
             return
 
-    async def astream(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, **kwargs) -> AsyncGenerator[str, None]:
+    async def astream(self, prompt: str, tools: Optional[List[Dict[str, Any]]] = None, session_id: Optional[Union[str, Session]] = None, **kwargs) -> AsyncGenerator[str, None]:
         """
         Execute an asynchronous streaming completion with automatic tool calling.
         """
-        messages = self._prepare_messages(prompt)
+        actual_session_id = session_id.id if isinstance(session_id, Session) else (session_id or str(uuid.uuid4()))
+        messages = self._prepare_messages(prompt, actual_session_id=actual_session_id)
         tools = tools or self.tools
+        new_turns = []
         
         while True:
             response = await litellm.acompletion(
@@ -460,7 +539,7 @@ class LiteLLMAgent(BaseAgent):
                                     if last_tc.function.arguments is None:
                                         last_tc.function.arguments = ""
                                     last_tc.function.arguments += tc_delta.function.arguments
-
+  
             tool_calls = []
             for idx in sorted(tool_calls_by_index.keys()):
                 for tc_obj in tool_calls_by_index[idx]:
@@ -473,24 +552,30 @@ class LiteLLMAgent(BaseAgent):
                                 "arguments": tc_obj.function.arguments
                             }
                         })
-
+  
             if tool_calls:
                 if self._should_handle_sequentially():
                     tool_calls = [tool_calls[0]]
-
-                assistant_msg = {"role": "assistant", "tool_calls": tool_calls, "content": full_content or None}
+  
+                assistant_msg = self._sanitize_message({"role": "assistant", "tool_calls": tool_calls, "content": full_content})
                 messages.append(assistant_msg)
+                new_turns.append(assistant_msg)
                 
-                for tool_call in tool_calls:
-                    result = self._execute_tool(tool_call)
-                    messages.append({
-                        "role": "tool",
-                        "tool_call_id": tool_call["id"],
-                        "name": tool_call["function"]["name"],
-                        "content": str(result)
-                    })
+                if self._should_handle_sequentially():
+                    for tool_call in tool_calls:
+                        result = await self._aexecute_tool(tool_call)
+                        messages.append(result)
+                        new_turns.append(result)
+                else:
+                    import asyncio
+                    results = await asyncio.gather(*[self._aexecute_tool(tc) for tc in tool_calls])
+                    for res in results:
+                        messages.append(res)
+                        new_turns.append(res)
                 continue
             
-            messages.append({"role": "assistant", "content": full_content})
-            self._update_history(messages)
+            final_msg = self._sanitize_message({"role": "assistant", "content": full_content})
+            messages.append(final_msg)
+            new_turns.append(final_msg)
+            self._update_history(new_turns, actual_session_id=actual_session_id)
             return

litellm_adk/core/context.py

@@ -0,0 +1,84 @@
+from typing import List, Dict, Any, Optional
+import litellm
+from ..observability.logger import adk_logger
+
+class ContextManager:
+    """
+    Handles token counting and history truncation/summarization.
+    """
+    
+    @staticmethod
+    def count_tokens(messages: List[Dict[str, Any]], model: str) -> int:
+        """
+        Calculate the number of tokens in a list of messages.
+        Uses cached 'token_count' if available, otherwise LiteLLM's token_counter.
+        """
+        # optimization: if it's a single message with a cached count, use it
+        if len(messages) == 1 and "token_count" in messages[0]:
+            return messages[0]["token_count"]
+            
+        try:
+            # Pass a copy to avoid any potential in-place modifications by token_counter
+            return litellm.token_counter(model=model, messages=[m.copy() for m in messages])
+        except Exception as e:
+            adk_logger.warning(f"Token counting failed for model {model}: {e}. Falling back to estimate.")
+            # Rough estimate: 4 chars per token
+            return sum(len(str(m.get("content", ""))) for m in messages) // 4
+
+    @staticmethod
+    def truncate_history(
+        messages: List[Dict[str, Any]], 
+        model: str, 
+        max_tokens: int,
+        reserve_tokens: int = 500
+    ) -> List[Dict[str, Any]]:
+        """
+        Truncate history to fit within max_tokens, always preserving the system prompt
+        and the latest message.
+        """
+        if not messages:
+            return []
+            
+        # 1. Separate System Prompt
+        system_prompt = None
+        if messages[0].get("role") == "system":
+            system_prompt = messages[0]
+            other_messages = messages[1:]
+        else:
+            other_messages = messages
+
+        # 2. Calculate Budget
+        actual_reserve = min(reserve_tokens, int(max_tokens * 0.2))
+        allowed_tokens = max_tokens - actual_reserve
+        
+        if system_prompt:
+            allowed_tokens -= ContextManager.count_tokens([system_prompt], model)
+
+        # 3. Quick Check: Is truncation even needed?
+        # This one call avoids the N calls in the loop below for most turns
+        if ContextManager.count_tokens(other_messages, model) <= allowed_tokens:
+            return messages
+
+        # 4. Truncate (Keeping the LATEST messages)
+        truncated = []
+        current_tokens = 0
+        
+        if other_messages:
+            last_msg = other_messages[-1]
+            truncated.append(last_msg)
+            current_tokens += ContextManager.count_tokens([last_msg], model)
+            
+            for msg in reversed(other_messages[:-1]):
+                msg_tokens = ContextManager.count_tokens([msg], model)
+                if current_tokens + msg_tokens > allowed_tokens:
+                    break
+                truncated.insert(0, msg)
+                current_tokens += msg_tokens
+
+        # 5. Reconstruct
+        result = []
+        if system_prompt:
+            result.append(system_prompt)
+        result.extend(truncated)
+        
+        return result

litellm_adk/memory/base.py

@@ -25,3 +25,13 @@ class BaseMemory(ABC):
     def clear(self, session_id: str):
         """Clear history for a session."""
         pass
+
+    @abstractmethod
+    def get_session_metadata(self, session_id: str) -> Dict[str, Any]:
+        """Retrieve metadata/state for a given session."""
+        pass
+
+    @abstractmethod
+    def save_session_metadata(self, session_id: str, metadata: Dict[str, Any]):
+        """Save/Update metadata/state for a given session."""
+        pass

litellm_adk/memory/file.py

@@ -5,18 +5,24 @@ from .base import BaseMemory
 
 class FileMemory(BaseMemory):
     """
-    JSON file-based persistence for conversation history.
+    JSON file-based persistence for conversation history and session metadata.
     """
     def __init__(self, file_path: str = "conversations.json"):
         self.file_path = file_path
-        self._cache: Dict[str, List[Dict[str, Any]]] = {}
+        self._cache: Dict[str, Dict[str, Any]] = {}
         self._load()
 
     def _load(self):
         if os.path.exists(self.file_path):
             with open(self.file_path, "r", encoding="utf-8") as f:
                 try:
-                    self._cache = json.load(f)
+                    data = json.load(f)
+                    # Support legacy format (where values were just lists)
+                    for k, v in data.items():
+                        if isinstance(v, list):
+                            self._cache[k] = {"messages": v, "metadata": {}}
+                        else:
+                            self._cache[k] = v
                 except json.JSONDecodeError:
                     self._cache = {}
         else:
@@ -27,21 +33,30 @@ class FileMemory(BaseMemory):
             json.dump(self._cache, f, indent=2, ensure_ascii=False)
 
     def get_messages(self, session_id: str) -> List[Dict[str, Any]]:
-        return self._cache.get(session_id, []).copy()
+        return self._cache.get(session_id, {}).get("messages", []).copy()
 
     def add_message(self, session_id: str, message: Dict[str, Any]):
         if session_id not in self._cache:
-            self._cache[session_id] = []
-        self._cache[session_id].append(message)
+            self._cache[session_id] = {"messages": [], "metadata": {}}
+        self._cache[session_id]["messages"].append(message)
         self._save()
 
     def add_messages(self, session_id: str, messages: List[Dict[str, Any]]):
         if session_id not in self._cache:
-            self._cache[session_id] = []
-        self._cache[session_id].extend(messages)
+            self._cache[session_id] = {"messages": [], "metadata": {}}
+        self._cache[session_id]["messages"].extend(messages)
         self._save()
 
     def clear(self, session_id: str):
         if session_id in self._cache:
-            self._cache[session_id] = []
+            self._cache[session_id] = {"messages": [], "metadata": {}}
             self._save()
+
+    def get_session_metadata(self, session_id: str) -> Dict[str, Any]:
+        return self._cache.get(session_id, {}).get("metadata", {}).copy()
+
+    def save_session_metadata(self, session_id: str, metadata: Dict[str, Any]):
+        if session_id not in self._cache:
+            self._cache[session_id] = {"messages": [], "metadata": {}}
+        self._cache[session_id]["metadata"] = metadata
+        self._save()

litellm_adk/memory/in_memory.py

@@ -3,10 +3,11 @@ from .base import BaseMemory
 
 class InMemoryMemory(BaseMemory):
     """
-    Standard in-memory store for conversation history.
+    Standard in-memory store for conversation history and session metadata.
     """
     def __init__(self):
         self._storage: Dict[str, List[Dict[str, Any]]] = {}
+        self._metadata: Dict[str, Dict[str, Any]] = {}
 
     def get_messages(self, session_id: str) -> List[Dict[str, Any]]:
         return self._storage.get(session_id, []).copy()
@@ -24,3 +25,11 @@ class InMemoryMemory(BaseMemory):
     def clear(self, session_id: str):
         if session_id in self._storage:
             self._storage[session_id] = []
+        if session_id in self._metadata:
+            self._metadata[session_id] = {}
+
+    def get_session_metadata(self, session_id: str) -> Dict[str, Any]:
+        return self._metadata.get(session_id, {}).copy()
+
+    def save_session_metadata(self, session_id: str, metadata: Dict[str, Any]):
+        self._metadata[session_id] = metadata.copy()

litellm_adk/memory/mongodb.py

@@ -4,7 +4,7 @@ import pymongo
 
 class MongoDBMemory(BaseMemory):
     """
-    MongoDB-based persistence for conversation history.
+    MongoDB-based persistence for conversation history and session metadata.
     """
     def __init__(
         self, 
@@ -21,8 +21,6 @@ class MongoDBMemory(BaseMemory):
     def get_messages(self, session_id: str) -> List[Dict[str, Any]]:
         doc = self.collection.find_one({"session_id": session_id})
         if doc:
-            # MongoDB returns a list of dicts, but we need to ensure 
-            # we return a copy to prevent in-place modifications
             return list(doc.get("messages", []))
         return []
 
@@ -43,5 +41,18 @@ class MongoDBMemory(BaseMemory):
     def clear(self, session_id: str):
         self.collection.update_one(
             {"session_id": session_id},
-            {"$set": {"messages": []}}
+            {"$set": {"messages": [], "metadata": {}}}
+        )
+
+    def get_session_metadata(self, session_id: str) -> Dict[str, Any]:
+        doc = self.collection.find_one({"session_id": session_id})
+        if doc:
+            return doc.get("metadata", {})
+        return {}
+
+    def save_session_metadata(self, session_id: str, metadata: Dict[str, Any]):
+        self.collection.update_one(
+            {"session_id": session_id},
+            {"$set": {"metadata": metadata}},
+            upsert=True
         )

litellm_adk/session/__init__.py

@@ -0,0 +1,3 @@
+from .session import Session
+
+__all__ = ["Session"]

litellm_adk/session/session.py

@@ -0,0 +1,24 @@
+from typing import Dict, Any, Optional
+from pydantic import BaseModel, Field
+import uuid
+import time
+
+class Session(BaseModel):
+    """
+    Represents a conversation session with metadata and state.
+    """
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    user_id: Optional[str] = None
+    app_name: Optional[str] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    state: Dict[str, Any] = Field(default_factory=dict)
+    created_at: float = Field(default_factory=time.time)
+    updated_at: float = Field(default_factory=time.time)
+
+    def update_state(self, key: str, value: Any):
+        self.state[key] = value
+        self.updated_at = time.time()
+
+    def update_metadata(self, key: str, value: Any):
+        self.metadata[key] = value
+        self.updated_at = time.time()

litellm_adk/tools/registry.py

@@ -79,7 +79,37 @@ class ToolRegistry:
             raise ValueError(f"Tool '{name}' not found in registry.")
             
         adk_logger.info(f"Executing tool: {name} with args: {kwargs}")
-        return self._tools[name]["func"](**kwargs)
+        func = self._tools[name]["func"]
+        
+        # Handle both sync and async functions if called synchronously
+        if inspect.iscoroutinefunction(func):
+            import asyncio
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    # If already in a loop, we can't block. This is a fallback risk.
+                    adk_logger.warning(f"Sync execution of async tool '{name}' in running loop.")
+                    return asyncio.run_coroutine_threadsafe(func(**kwargs), loop).result()
+                return asyncio.run(func(**kwargs))
+            except RuntimeError:
+                return asyncio.run(func(**kwargs))
+        
+        return func(**kwargs)
+
+    async def aexecute(self, name: str, **kwargs) -> Any:
+        """
+        Asynchronously executes a registered tool by name.
+        """
+        if name not in self._tools:
+            raise ValueError(f"Tool '{name}' not found in registry.")
+            
+        adk_logger.info(f"Executing tool (async): {name} with args: {kwargs}")
+        func = self._tools[name]["func"]
+        
+        if inspect.iscoroutinefunction(func):
+            return await func(**kwargs)
+        else:
+            return func(**kwargs)
 
 # Global tool registry
 tool_registry = ToolRegistry()

{litellm_adk-0.2.0.dist-info → litellm_adk-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: litellm-adk
-Version: 0.2.0
+Version: 0.2.1
 Summary: Production-grade multiservice Agent Development Kit
 License-Expression: MIT
 Classifier: Programming Language :: Python :: 3

litellm_adk-0.2.1.dist-info/RECORD

@@ -0,0 +1,21 @@
+litellm_adk/__init__.py,sha256=k82anVFffmjIP32AcLCFsEtKjk6YXbEpZ8spy1UWFko,414
+litellm_adk/agents/__init__.py,sha256=KsiCcyYyn0iJo0sZsd7n3nJ5ezVEqJqrrTP1b9ryG0M,69
+litellm_adk/config/settings.py,sha256=sdI4PrJKzIRJPm5vEBKdiecRcjrHNASsYbmtOiBnY_c,1182
+litellm_adk/core/agent.py,sha256=bCigFJCKAkbAeiMt68mbDSq-zoSlYCQgMp3MowksKok,26856
+litellm_adk/core/base.py,sha256=ov2bZk_a15FFGsQSdKwHrQ1cvALdZM8ByK5hGvFWyL0,386
+litellm_adk/core/context.py,sha256=7tPPRAQ79EumA-mnaG39bsG8PpLqwCpBT-f8zq0Ftn0,3162
+litellm_adk/memory/__init__.py,sha256=ICPUbV0PsTHEQSm0S35_d1ToeyrgMVFs_hRokvRRJL4,212
+litellm_adk/memory/base.py,sha256=gfwhcORWZJl4qqb0u2uIPOLos1MYwjwkqiLT242ExSM,1138
+litellm_adk/memory/file.py,sha256=_9oFobFE3tbM1f7L3jBBEEVXl6SM4mMITnNxzTO5iKg,2519
+litellm_adk/memory/in_memory.py,sha256=K5sR_z5W8Yyd8Uj3DDD6an6RggBLXiYAjYVJgG7z94o,1384
+litellm_adk/memory/mongodb.py,sha256=Sxsuemfxhxm3NeyoYDWC6xnH56uZ2idSAViEGKK1cjo,2012
+litellm_adk/observability/logger.py,sha256=PXr20D7gtDIrg6eZD8Hm1-tfAuTXyUVDUMD9-8Aw32E,619
+litellm_adk/session/__init__.py,sha256=G1F-Xgmj-Aezewv977vZKflhmZnqEwrYkcwqMTyVpWI,55
+litellm_adk/session/session.py,sha256=CE-bovngwsRMjjYCC1gdDjJKnewRPJO8Qw2akMJTdbw,840
+litellm_adk/tools/__init__.py,sha256=J-Rkx-psP5sZXgcy5h4mygvQd-tZUONKLYt4LSOiEV8,82
+litellm_adk/tools/registry.py,sha256=UTZrWuSWDhdu8arZt9bfkHz8V0sy-Odpiki6Sh6Up4E,4337
+litellm_adk-0.2.1.dist-info/licenses/LICENSE,sha256=BfYjX2LxngGX9t6Dk1Y5ptJNAkKcQuGG-OAR9jsKUGM,1091
+litellm_adk-0.2.1.dist-info/METADATA,sha256=PZFvhminGSfIw_pOMpFYcoRb6c3fIONvjFzcW6MYAjg,2981
+litellm_adk-0.2.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+litellm_adk-0.2.1.dist-info/top_level.txt,sha256=30MPgkTEjMUe8z-jnjMM2vbtqdghK_isd_ufRQ1w2hM,12
+litellm_adk-0.2.1.dist-info/RECORD,,

litellm_adk-0.2.0.dist-info/RECORD

@@ -1,18 +0,0 @@
-litellm_adk/__init__.py,sha256=DfQsjG9D_5OvxX-8tOtGY3kuCxYb6Hn_I-2BvnTqf0M,368
-litellm_adk/agents/__init__.py,sha256=KsiCcyYyn0iJo0sZsd7n3nJ5ezVEqJqrrTP1b9ryG0M,69
-litellm_adk/config/settings.py,sha256=sdI4PrJKzIRJPm5vEBKdiecRcjrHNASsYbmtOiBnY_c,1182
-litellm_adk/core/agent.py,sha256=XE6dJZ6q3ZDnaMLtYHislnx6-IDTlXrrBx5q1_4Nt-4,22248
-litellm_adk/core/base.py,sha256=ov2bZk_a15FFGsQSdKwHrQ1cvALdZM8ByK5hGvFWyL0,386
-litellm_adk/memory/__init__.py,sha256=ICPUbV0PsTHEQSm0S35_d1ToeyrgMVFs_hRokvRRJL4,212
-litellm_adk/memory/base.py,sha256=Bm33oPaLNOJdG0RJGc38g387GSMyi_ymQjOMlDexTyk,788
-litellm_adk/memory/file.py,sha256=C0pB1pWJ4HtjCn6ICe54pL3cVmVnPI6D5jRvQaffgE4,1623
-litellm_adk/memory/in_memory.py,sha256=AVMV7iqb-UvbPE-CZmRi14LkV-7hEqrkwtkrwlxvy_w,951
-litellm_adk/memory/mongodb.py,sha256=M7IQsgahT6ALSNTQ2AKjSUWGR7uuz-KielBYIu_oLVk,1657
-litellm_adk/observability/logger.py,sha256=PXr20D7gtDIrg6eZD8Hm1-tfAuTXyUVDUMD9-8Aw32E,619
-litellm_adk/tools/__init__.py,sha256=J-Rkx-psP5sZXgcy5h4mygvQd-tZUONKLYt4LSOiEV8,82
-litellm_adk/tools/registry.py,sha256=M_48BpN0XSea_3msjGSyyDDRWu9uBNDLDtLh9Vh5yp8,3089
-litellm_adk-0.2.0.dist-info/licenses/LICENSE,sha256=BfYjX2LxngGX9t6Dk1Y5ptJNAkKcQuGG-OAR9jsKUGM,1091
-litellm_adk-0.2.0.dist-info/METADATA,sha256=x7FfuvUhi-u365L7AUzXP3qfF_9PQQ-dNwL7L2iNd14,2981
-litellm_adk-0.2.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-litellm_adk-0.2.0.dist-info/top_level.txt,sha256=30MPgkTEjMUe8z-jnjMM2vbtqdghK_isd_ufRQ1w2hM,12
-litellm_adk-0.2.0.dist-info/RECORD,,