adaptive-memory-multi-model-router 2.15.3 → 2.15.5

package/adapters/a3m_adapter/adapter/langgraph.py

@@ -0,0 +1,196 @@
+"""
+A3M Router Adapter for LangGraph (Microsoft's agent framework).
+
+Drop-in replacement for LangGraph's stateful agent that routes through A3M Router
+for intelligent, cost-optimized multi-step conversations.
+
+Usage:
+    from langgraph.prebuilt import create_react_agent
+    from a3m_adapter import A3MLangGraphAdapter
+    
+    adapter = A3MLangGraphAdapter(model='auto', parallel_ensemble=2)
+    
+    agent = create_react_agent(adapter, tools=[...])
+    
+    result = agent.invoke({"messages": [{"role": "user", "content": "Hello"]})
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, TypedDict
+
+logger = logging.getLogger(__name__)
+
+LANGGRAPH_AVAILABLE = False
+try:
+    import langgraph
+    from langgraph.prebuilt import create_react_agent
+    from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
+    LANGGRAPH_AVAILABLE = True
+except ImportError:
+    logger.warning("LangGraph not installed. Install with: pip install langgraph")
+
+A3M_AVAILABLE = False
+try:
+    from a3m.router import A3MRouter, RouteResponse
+    A3M_AVAILABLE = True
+except ImportError:
+    logger.warning(
+        "A3M Router not installed. Install with: pip install adaptive-memory-multi-model-router"
+    )
+
+
+class A3MLangGraphAdapter:
+    """
+    A3M Router adapter for LangGraph's prebuilt agents.
+
+    Enables LangGraph agents to use A3M Router for automatic model selection
+    across 47+ providers with cost optimization and stateful conversations.
+    """
+
+    def __init__(
+        self,
+        model: str = "auto",
+        temperature: float = 0.7,
+        max_tokens: int = 4096,
+        parallel_ensemble: int = 1,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize A3M Router adapter for LangGraph.
+        """
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.parallel_ensemble = parallel_ensemble
+        self.api_key = api_key
+        self._a3m_router = None
+        self._initialized = False
+        self._kwargs = kwargs
+
+    def _ensure_router(self) -> None:
+        """Lazily initialize the A3M router."""
+        if self._initialized:
+            return
+        
+        if not A3M_AVAILABLE:
+            raise ImportError(
+                "A3M Router is not installed. "
+                "Install with: pip install adaptive-memory-multi-model-router"
+            )
+        
+        self._a3m_router = A3MRouter(
+            model=self.model,
+            temperature=self.temperature,
+            parallel_ensemble=self.parallel_ensemble,
+        )
+        self._initialized = True
+        logger.info(
+            "A3M Router initialized for LangGraph: model=%s, ensemble=%d",
+            self.model,
+            self.parallel_ensemble,
+        )
+
+    def get_model(self):
+        """
+        Get the underlying model for LangGraph.
+        
+        Returns an object compatible with LangGraph's prebuilt agents.
+        """
+        self._ensure_router()
+        return self
+
+    def __call__(
+        self,
+        state: Dict[str, Any],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        LangGraph-compatible callable for node execution.
+        
+        Args:
+            state: LangGraph state dict with 'messages' key
+            
+        Returns:
+            Updated state dict
+        """
+        self._ensure_router()
+        
+        messages = state.get("messages", [])
+        
+        # Convert LangGraph messages to A3M format
+        a3m_messages = self._convert_messages(messages)
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=a3m_messages,
+                temperature=kwargs.get("temperature", self.temperature),
+                max_tokens=kwargs.get("max_tokens", self.max_tokens),
+                **kwargs,
+            ),
+        )
+        
+        # Add response to messages
+        new_messages = messages + [
+            AIMessage(content=route_result.content)
+        ]
+        
+        return {
+            **state,
+            "messages": new_messages,
+        }
+
+    def _convert_messages(
+        self,
+        messages: List[BaseMessage],
+    ) -> List[Dict[str, Any]]:
+        """Convert LangGraph messages to A3M format."""
+        a3m_messages = []
+        for msg in messages:
+            if isinstance(msg, HumanMessage):
+                a3m_messages.append({"role": "user", "content": msg.content})
+            elif isinstance(msg, AIMessage):
+                a3m_messages.append({"role": "assistant", "content": msg.content})
+            else:
+                a3m_messages.append({"role": "user", "content": str(msg)})
+        return a3m_messages
+
+    async def ainvoke(
+        self,
+        state: Dict[str, Any],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Async version of __call__."""
+        self._ensure_router()
+        
+        messages = state.get("messages", [])
+        a3m_messages = self._convert_messages(messages)
+        
+        route_result = await self._a3m_router.aroute(
+            messages=a3m_messages,
+            temperature=kwargs.get("temperature", self.temperature),
+            max_tokens=kwargs.get("max_tokens", self.max_tokens),
+            **kwargs,
+        )
+        
+        new_messages = messages + [
+            AIMessage(content=route_result.content)
+        ]
+        
+        return {
+            **state,
+            "messages": new_messages,
+        }
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MLangGraphAdapter("
+            f"model={self.model!r}, "
+            f"temperature={self.temperature}, "
+            f"ensemble={self.parallel_ensemble})"
+        )

package/adapters/a3m_adapter/adapter/llamaindex.py

@@ -0,0 +1,162 @@
+"""
+A3M Router Adapter for LlamaIndex.
+
+Drop-in replacement for LlamaIndex's BaseLLM that routes through A3M Router
+for intelligent, cost-optimized model selection.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional, Sequence
+
+logger = logging.getLogger(__name__)
+
+# Check availability
+LLAMAINDEX_AVAILABLE = False
+_llama_metadata_class = None
+try:
+    from llama_index.core.base.llms.base import BaseLLM, CompletionResponse
+    from llama_index.core.base.llms.types import ChatMessage
+    LLAMAINDEX_AVAILABLE = True
+    try:
+        from llama_index.core.base.llms.base import LLMMetadata
+        _llama_metadata_class = LLMMetadata
+    except ImportError:
+        pass
+except ImportError:
+    logger.warning("LlamaIndex not installed. Install with: pip install llama-index")
+
+A3M_AVAILABLE = False
+try:
+    from a3m.router import A3MRouter, RouteResponse
+    A3M_AVAILABLE = True
+except ImportError:
+    logger.warning("A3M Router not installed. Install with: pip install adaptive-memory-multi-model-router")
+
+
+class A3MLlamaIndexAdapter:
+    """
+    A3M Router adapter for LlamaIndex's BaseLLM interface.
+
+    Routes prompts through A3M Router to automatically select the cheapest
+    capable model across 47+ LLM providers.
+    """
+
+    def __init__(
+        self,
+        model: str = "auto",
+        temperature: float = 0.0,
+        max_tokens: Optional[int] = 4096,
+        parallel_ensemble: int = 1,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize A3M Router adapter.
+        """
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.parallel_ensemble = parallel_ensemble
+        self.api_key = api_key
+        self._a3m_router = None
+        self._initialized = False
+
+    def _ensure_router(self) -> None:
+        """Lazily initialize the A3M router."""
+        if self._initialized:
+            return
+        
+        if not A3M_AVAILABLE:
+            raise ImportError(
+                "A3M Router is not installed. "
+                "Install with: pip install adaptive-memory-multi-model-router"
+            )
+        
+        self._a3m_router = A3MRouter(
+            model=self.model,
+            temperature=self.temperature,
+            parallel_ensemble=self.parallel_ensemble,
+        )
+        self._initialized = True
+        logger.info(
+            "A3M Router initialized: model=%s, ensemble=%d",
+            self.model,
+            self.parallel_ensemble,
+        )
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        """Return LLM metadata as a dict (framework-agnostic)."""
+        return {
+            "context_window": 128000,
+            "num_output": self.max_tokens or 4096,
+            "model_name": self.model,
+            "is_chat_model": True,
+        }
+
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        """Complete a prompt using A3M Router."""
+        self._ensure_router()
+        
+        messages = [{"role": "user", "content": prompt}]
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                **kwargs,
+            ),
+        )
+        
+        return CompletionResponse(text=route_result.content, raw=route_result)
+
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> CompletionResponse:
+        """Chat completion using A3M Router."""
+        self._ensure_router()
+        
+        a3m_messages = self._convert_messages(messages)
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=a3m_messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                **kwargs,
+            ),
+        )
+        
+        return CompletionResponse(text=route_result.content, raw=route_result)
+
+    def _convert_messages(self, messages: Sequence[ChatMessage]) -> List[Dict[str, Any]]:
+        """Convert LlamaIndex ChatMessages to A3M format."""
+        a3m_messages = []
+        for msg in messages:
+            role = msg.role.value if hasattr(msg.role, 'value') else str(msg.role).lower()
+            role_map = {
+                "system": "system",
+                "user": "user", 
+                "assistant": "assistant",
+                "tool": "tool",
+                "function": "function",
+            }
+            a3m_role = role_map.get(role, "user")
+            a3m_messages.append({"role": a3m_role, "content": msg.content})
+        return a3m_messages
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MLlamaIndexAdapter("
+            f"model={self.model!r}, "
+            f"temperature={self.temperature}, "
+            f"max_tokens={self.max_tokens}, "
+            f"ensemble={self.parallel_ensemble})"
+        )

package/adapters/a3m_adapter/adapter/pinecone.py

@@ -0,0 +1,217 @@
+"""
+A3M Router Adapter for Pinecone Vector Database.
+
+Enables Pinecone's managed vector database to use A3M Router for
+intelligent query routing and cost-optimized embeddings.
+
+Usage:
+    from pinecone import Pinecone
+    from a3m_adapter import A3MPineconeAdapter
+    
+    # Create A3M-powered embeddings
+    embed_adapter = A3MPineconeAdapter(model="auto")
+    
+    # Generate embeddings
+    embedding = embed_adapter.embed("What is quantum computing?")
+    
+    # Query Pinecone
+    results = index.query(
+        vector=embedding,
+        top_k=5,
+    )
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+A3M_AVAILABLE = False
+try:
+    from a3m.router import A3MRouter
+    A3M_AVAILABLE = True
+except ImportError:
+    logger.warning(
+        "A3M Router not installed. Install with: pip install adaptive-memory-multi-model-router"
+    )
+
+
+class A3MPineconeAdapter:
+    """
+    A3M Router adapter for Pinecone embeddings.
+
+    Provides intelligent embedding generation through A3M Router
+    with automatic model selection for cost optimization.
+    """
+
+    def __init__(
+        self,
+        model: str = "auto",
+        embed_model: str = "auto",
+        parallel_ensemble: int = 1,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize A3M Router adapter for Pinecone.
+        """
+        self.model = model
+        self.embed_model = embed_model or "auto"
+        self.parallel_ensemble = parallel_ensemble
+        self.api_key = api_key
+        self._a3m_router = None
+        self._initialized = False
+        self._kwargs = kwargs
+
+    def _ensure_router(self) -> None:
+        """Lazily initialize the A3M router."""
+        if self._initialized:
+            return
+        
+        if not A3M_AVAILABLE:
+            raise ImportError(
+                "A3M Router is not installed. "
+                "Install with: pip install adaptive-memory-multi-model-router"
+            )
+        
+        self._a3m_router = A3MRouter(
+            model=self.model,
+            parallel_ensemble=self.parallel_ensemble,
+        )
+        self._initialized = True
+        logger.info(
+            "A3M Router initialized for Pinecone: embed_model=%s",
+            self.embed_model,
+        )
+
+    def embed(
+        self,
+        texts: List[str],
+        **kwargs: Any,
+    ) -> List[List[float]]:
+        """
+        Generate embeddings for texts using A3M Router.
+        
+        Args:
+            texts: List of text strings to embed
+            
+        Returns:
+            List of embedding vectors
+        """
+        self._ensure_router()
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        
+        # For embeddings, we typically call the router with a special embedding mode
+        # Since A3M supports /v1/embeddings endpoint
+        results = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.embed(
+                texts=texts,
+                model=self.embed_model,
+                **kwargs,
+            ),
+        )
+        
+        return results
+
+    def embed_query(
+        self,
+        text: str,
+        **kwargs: Any,
+    ) -> List[float]:
+        """
+        Generate embedding for a single query.
+        
+        Args:
+            text: Text to embed
+            
+        Returns:
+            Embedding vector
+        """
+        embeddings = self.embed([text], **kwargs)
+        return embeddings[0] if embeddings else []
+
+    async def aembed(
+        self,
+        texts: List[str],
+        **kwargs: Any,
+    ) -> List[List[float]]:
+        """Async version of embed."""
+        self._ensure_router()
+        
+        results = await self._a3m_router.aembed(
+            texts=texts,
+            model=self.embed_model,
+            **kwargs,
+        )
+        
+        return results
+
+    def rag_query(
+        self,
+        query: str,
+        index,
+        top_k: int = 5,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Perform RAG query: embed + Pinecone search + context.
+        
+        Args:
+            query: The search query
+            index: Pinecone index to query
+            top_k: Number of results to retrieve
+            
+        Returns:
+            Dict with 'results', 'context', 'provider', 'cost'
+        """
+        # 1. Embed query
+        query_embedding = self.embed_query(query)
+        
+        # 2. Search Pinecone
+        search_results = index.query(
+            vector=query_embedding,
+            top_k=top_k,
+            include_metadata=True,
+        )
+        
+        # 3. Build context from results
+        context = "\n\n".join([
+            match.get('metadata', {}).get('text', str(match.get('id', '')))
+            for match in search_results.get('matches', [])[:3]
+        ])
+        
+        # 4. Route the full query through A3M
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=[{
+                    "role": "user",
+                    "content": f"Context:\n{context}\n\nQuestion: {query}"
+                }],
+                temperature=self._kwargs.get("temperature", 0.7),
+                **kwargs,
+            ),
+        )
+        
+        return {
+            "results": search_results.get('matches', []),
+            "context": context,
+            "answer": route_result.content,
+            "provider": getattr(route_result, 'provider', 'a3m'),
+            "cost": getattr(route_result, 'cost', 0.0),
+        }
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MPineconeAdapter("
+            f"model={self.model!r}, "
+            f"embed_model={self.embed_model!r}, "
+            f"ensemble={self.parallel_ensemble})"
+        )