adaptive-memory-multi-model-router 2.16.0 → 2.16.2

package/adapters/a3m_adapter/adapter/haystack.py

@@ -0,0 +1,197 @@
+"""
+A3M Router Adapter for Haystack (Deepset's RAG framework).
+
+Drop-in replacement for Haystack's OpenAIGenerator that routes through A3M Router
+for intelligent, cost-optimized RAG pipelines.
+
+Usage:
+    from haystack import Pipeline
+    from haystack.nodes import Retriever, PromptNode
+    from a3m_adapter import A3MHaystackAdapter
+    
+    prompt_node = PromptNode(
+        "auto",
+        api_key=None,
+        generator_type='openai',
+        model_adapter=A3MHaystackAdapter(model='auto', parallel_ensemble=2),
+    )
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+HAYSTACK_AVAILABLE = False
+try:
+    from haystack.nodes.base import BaseGenerator
+    HAYSTACK_AVAILABLE = True
+except ImportError:
+    logger.warning("Haystack not installed. Install with: pip install farm-haystack")
+
+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 A3MHaystackAdapter:
+    """
+    A3M Router adapter for Haystack's PromptNode.
+
+    Enables Haystack RAG pipelines to use A3M Router for automatic model selection
+    across 47+ providers with cost optimization.
+    """
+
+    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 Haystack.
+        """
+        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 Haystack: model=%s",
+            self.model,
+        )
+
+    def predict(
+        self,
+        query: str,
+        documents: Optional[List[Any]] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Generate answer from query and optional retrieved documents.
+        
+        Args:
+            query: The search query
+            documents: Optional list of retrieved documents for RAG
+            
+        Returns:
+            Dict with 'answers', 'provider', 'cost'
+        """
+        self._ensure_router()
+        
+        # Build context from documents if provided
+        if documents:
+            context = "\n\n".join([
+                f"Document {i+1}: {getattr(doc, 'content', str(doc))}"
+                for i, doc in enumerate(documents[:5])  # Limit to 5 docs
+            ])
+            prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"
+        else:
+            prompt = query
+        
+        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=kwargs.get("temperature", self.temperature),
+                max_tokens=kwargs.get("max_tokens", self.max_tokens),
+                **kwargs,
+            ),
+        )
+        
+        return {
+            "answers": [{"answer": route_result.content, "score": 1.0}],
+            "provider": getattr(route_result, 'provider', 'a3m'),
+            "cost": getattr(route_result, 'cost', 0.0),
+        }
+
+    async def apredict(
+        self,
+        query: str,
+        documents: Optional[List[Any]] = None,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Async predict for Haystack."""
+        self._ensure_router()
+        
+        if documents:
+            context = "\n\n".join([
+                f"Document {i+1}: {getattr(doc, 'content', str(doc))}"
+                for i, doc in enumerate(documents[:5])
+            ])
+            prompt = f"Context:\n{context}\n\nQuestion: {query}\n\nAnswer:"
+        else:
+            prompt = query
+        
+        messages = [{"role": "user", "content": prompt}]
+        
+        route_result = await self._a3m_router.aroute(
+            messages=messages,
+            temperature=kwargs.get("temperature", self.temperature),
+            max_tokens=kwargs.get("max_tokens", self.max_tokens),
+            **kwargs,
+        )
+        
+        return {
+            "answers": [{"answer": route_result.content, "score": 1.0}],
+            "provider": getattr(route_result, 'provider', 'a3m'),
+            "cost": getattr(route_result, 'cost', 0.0),
+        }
+
+    def run(
+        self,
+        query: str,
+        documents: Optional[List[Any]] = None,
+        **kwargs: Any,
+    ) -> tuple[Dict[str, Any], str]:
+        """
+        Haystack-compatible run method.
+        
+        Returns:
+            Tuple of (results dict, pipeline run metadata)
+        """
+        result = self.predict(query, documents, **kwargs)
+        return (result, "a3m-haystack")
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MHaystackAdapter("
+            f"model={self.model!r}, "
+            f"temperature={self.temperature}, "
+            f"max_tokens={self.max_tokens})"
+        )

package/adapters/a3m_adapter/adapter/langchain.py

@@ -0,0 +1,155 @@
+"""
+A3M Router Adapter for LangChain.
+
+Drop-in replacement for LangChain's ChatOpenAI that routes through A3M Router
+for intelligent, cost-optimized model selection across 47+ providers.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Check availability
+LANGCHAIN_AVAILABLE = False
+try:
+    from langchain_core.language_models import BaseChatModel
+    from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
+    from langchain_core.outputs import ChatGeneration, ChatResult, LLMResult
+    LANGCHAIN_AVAILABLE = True
+except ImportError:
+    logger.warning("LangChain not installed. Install with: pip install langchain langchain-core")
+
+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 A3MLangChainAdapter:
+    """
+    A3M Router adapter for LangChain's ChatOpenAI 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.
+
+        Args:
+            model: Model name or "auto" for automatic routing
+            temperature: Sampling temperature
+            max_tokens: Maximum tokens to generate
+            parallel_ensemble: Number of providers to run in parallel
+            api_key: A3M API key (optional)
+        """
+        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 _llm_type(self) -> str:
+        return "a3m_router"
+
+    def _generate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Any = None,
+        **kwargs: Any,
+    ) -> LLMResult:
+        """Generate a response using A3M Router."""
+        self._ensure_router()
+        
+        # Convert messages
+        a3m_messages = self._convert_messages(messages)
+
+        # Route through A3M
+        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,
+                stop=stop,
+                **kwargs,
+            ),
+        )
+
+        ai_message = AIMessage(content=route_result.content)
+        generation = ChatGeneration(message=ai_message)
+        return LLMResult(generations=[[generation]])
+
+    def _convert_messages(self, messages: List[BaseMessage]) -> List[Dict[str, Any]]:
+        """Convert LangChain messages to A3M format."""
+        a3m_messages = []
+        for msg in messages:
+            if isinstance(msg, SystemMessage):
+                a3m_messages.append({"role": "system", "content": msg.content})
+            elif isinstance(msg, HumanMessage):
+                a3m_messages.append({"role": "user", "content": msg.content})
+            elif isinstance(msg, AIMessage):
+                a3m_messages.append({"role": "assistant", "content": msg.content})
+            elif isinstance(msg, ToolMessage):
+                a3m_messages.append(
+                    {"role": "tool", "content": msg.content, "tool_call_id": msg.tool_call_id}
+                )
+            else:
+                a3m_messages.append({"role": "user", "content": str(msg)})
+        return a3m_messages
+
+    def bind_tools(self, tools: List[Dict[str, Any]], **kwargs: Any) -> "A3MLangChainAdapter":
+        """Bind tools for function calling."""
+        return self
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MLangChainAdapter("
+            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/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})"
+        )