a3m-router 2.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- a3m_router-2.2.0/PKG-INFO +98 -0
- a3m_router-2.2.0/README.md +55 -0
- a3m_router-2.2.0/a3m/__init__.py +35 -0
- a3m_router-2.2.0/a3m/adapters/__init__.py +21 -0
- a3m_router-2.2.0/a3m/adapters/langchain.py +190 -0
- a3m_router-2.2.0/a3m/adapters/llamaindex.py +249 -0
- a3m_router-2.2.0/a3m/adapters/qdrant.py +240 -0
- a3m_router-2.2.0/a3m/adapters/weaviate.py +263 -0
- a3m_router-2.2.0/a3m/client.py +195 -0
- a3m_router-2.2.0/a3m/models.py +40 -0
- a3m_router-2.2.0/a3m/sync_client.py +61 -0
- a3m_router-2.2.0/a3m_router.egg-info/PKG-INFO +98 -0
- a3m_router-2.2.0/a3m_router.egg-info/SOURCES.txt +17 -0
- a3m_router-2.2.0/a3m_router.egg-info/dependency_links.txt +1 -0
- a3m_router-2.2.0/a3m_router.egg-info/requires.txt +24 -0
- a3m_router-2.2.0/a3m_router.egg-info/top_level.txt +1 -0
- a3m_router-2.2.0/pyproject.toml +74 -0
- a3m_router-2.2.0/setup.cfg +4 -0
- a3m_router-2.2.0/setup.py +3 -0
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: a3m-router
|
|
3
|
+
Version: 2.2.0
|
|
4
|
+
Summary: Intelligent LLM routing for Python — routes to cheapest capable model across 47+ providers
|
|
5
|
+
License: MIT
|
|
6
|
+
Project-URL: Homepage, https://github.com/Das-rebel/a3m-router
|
|
7
|
+
Project-URL: Documentation, https://das-rebel.github.io/a3m-router
|
|
8
|
+
Project-URL: Repository, https://github.com/Das-rebel/a3m-router
|
|
9
|
+
Project-URL: Changelog, https://github.com/Das-rebel/a3m-router/releases
|
|
10
|
+
Keywords: llm,routing,openai,proxy,ai-gateway,routellm,litellm,multi-llm,model-selection,cost-optimization,langchain,llamaindex,rag,vector-search,embeddings,parallel,ensemble,router
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Operating System :: OS Independent
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
21
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
23
|
+
Requires-Python: >=3.8
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
Requires-Dist: httpx>=0.24.0
|
|
26
|
+
Provides-Extra: langchain
|
|
27
|
+
Requires-Dist: langchain>=0.1.0; extra == "langchain"
|
|
28
|
+
Provides-Extra: llamaindex
|
|
29
|
+
Requires-Dist: llamaindex>=0.9.0; extra == "llamaindex"
|
|
30
|
+
Provides-Extra: qdrant
|
|
31
|
+
Requires-Dist: qdrant-client>=1.7.0; extra == "qdrant"
|
|
32
|
+
Provides-Extra: weaviate
|
|
33
|
+
Requires-Dist: weaviate-client>=4.0.0; extra == "weaviate"
|
|
34
|
+
Provides-Extra: adapters
|
|
35
|
+
Requires-Dist: langchain>=0.1.0; extra == "adapters"
|
|
36
|
+
Requires-Dist: llamaindex>=0.9.0; extra == "adapters"
|
|
37
|
+
Requires-Dist: qdrant-client>=1.7.0; extra == "adapters"
|
|
38
|
+
Requires-Dist: weaviate-client>=4.0.0; extra == "adapters"
|
|
39
|
+
Provides-Extra: dev
|
|
40
|
+
Requires-Dist: pytest>=7.4.0; extra == "dev"
|
|
41
|
+
Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
|
|
42
|
+
Requires-Dist: black>=23.0.0; extra == "dev"
|
|
43
|
+
|
|
44
|
+
# A3M Router Python SDK
|
|
45
|
+
|
|
46
|
+
Intelligent LLM routing for Python — routes to cheapest capable model across 47+ providers.
|
|
47
|
+
|
|
48
|
+
## Installation
|
|
49
|
+
|
|
50
|
+
```bash
|
|
51
|
+
pip install a3m-router
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
## Usage
|
|
55
|
+
|
|
56
|
+
```python
|
|
57
|
+
from a3m import A3MRouter
|
|
58
|
+
|
|
59
|
+
router = A3MRouter() # localhost:8787
|
|
60
|
+
response = await router.chat("What is 2+2?")
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
## Framework Adapters
|
|
64
|
+
|
|
65
|
+
| Adapter | Import | Use Case |
|
|
66
|
+
|---------|--------|----------|
|
|
67
|
+
| **LangChain** | `from a3m import LangChainAdapter` | Chain-based AI workflows |
|
|
68
|
+
| **LlamaIndex** | `from a3m import LlamaIndexAdapter` | RAG and document QA |
|
|
69
|
+
| **Qdrant** | `from a3m import QdrantAdapter` | Vector search + RAG |
|
|
70
|
+
| **Weaviate** | `from a3m import WeaviateAdapter` | Vector search + RAG |
|
|
71
|
+
|
|
72
|
+
Install adapters:
|
|
73
|
+
```bash
|
|
74
|
+
pip install a3m-router[langchain] # LangChain
|
|
75
|
+
pip install a3m-router[llamaindex] # LlamaIndex
|
|
76
|
+
pip install a3m-router[qdrant] # Qdrant
|
|
77
|
+
pip install a3m-router[weaviate] # Weaviate
|
|
78
|
+
pip install a3m-router[all] # All adapters
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
## API Reference
|
|
82
|
+
|
|
83
|
+
| Method | Description |
|
|
84
|
+
|--------|-------------|
|
|
85
|
+
| `chat(message)` | Send chat message with auto-routing |
|
|
86
|
+
| `route(query)` | Get routing decision (no execution) |
|
|
87
|
+
| `stream_chat(message)` | Stream response tokens |
|
|
88
|
+
| `models()` | List available models |
|
|
89
|
+
| `health()` | Provider health status |
|
|
90
|
+
|
|
91
|
+
## Links
|
|
92
|
+
|
|
93
|
+
- [GitHub](https://github.com/Das-rebel/a3m-router)
|
|
94
|
+
- [npm Package](https://www.npmjs.com/package/adaptive-memory-multi-model-router)
|
|
95
|
+
|
|
96
|
+
## License
|
|
97
|
+
|
|
98
|
+
MIT
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
# A3M Router Python SDK
|
|
2
|
+
|
|
3
|
+
Intelligent LLM routing for Python — routes to cheapest capable model across 47+ providers.
|
|
4
|
+
|
|
5
|
+
## Installation
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install a3m-router
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Usage
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from a3m import A3MRouter
|
|
15
|
+
|
|
16
|
+
router = A3MRouter() # localhost:8787
|
|
17
|
+
response = await router.chat("What is 2+2?")
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
## Framework Adapters
|
|
21
|
+
|
|
22
|
+
| Adapter | Import | Use Case |
|
|
23
|
+
|---------|--------|----------|
|
|
24
|
+
| **LangChain** | `from a3m import LangChainAdapter` | Chain-based AI workflows |
|
|
25
|
+
| **LlamaIndex** | `from a3m import LlamaIndexAdapter` | RAG and document QA |
|
|
26
|
+
| **Qdrant** | `from a3m import QdrantAdapter` | Vector search + RAG |
|
|
27
|
+
| **Weaviate** | `from a3m import WeaviateAdapter` | Vector search + RAG |
|
|
28
|
+
|
|
29
|
+
Install adapters:
|
|
30
|
+
```bash
|
|
31
|
+
pip install a3m-router[langchain] # LangChain
|
|
32
|
+
pip install a3m-router[llamaindex] # LlamaIndex
|
|
33
|
+
pip install a3m-router[qdrant] # Qdrant
|
|
34
|
+
pip install a3m-router[weaviate] # Weaviate
|
|
35
|
+
pip install a3m-router[all] # All adapters
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
## API Reference
|
|
39
|
+
|
|
40
|
+
| Method | Description |
|
|
41
|
+
|--------|-------------|
|
|
42
|
+
| `chat(message)` | Send chat message with auto-routing |
|
|
43
|
+
| `route(query)` | Get routing decision (no execution) |
|
|
44
|
+
| `stream_chat(message)` | Stream response tokens |
|
|
45
|
+
| `models()` | List available models |
|
|
46
|
+
| `health()` | Provider health status |
|
|
47
|
+
|
|
48
|
+
## Links
|
|
49
|
+
|
|
50
|
+
- [GitHub](https://github.com/Das-rebel/a3m-router)
|
|
51
|
+
- [npm Package](https://www.npmjs.com/package/adaptive-memory-multi-model-router)
|
|
52
|
+
|
|
53
|
+
## License
|
|
54
|
+
|
|
55
|
+
MIT
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
"""A3M Router Python SDK"""
|
|
2
|
+
from .client import A3MRouter, A3MRouterError
|
|
3
|
+
from .sync_client import A3MRouterSync
|
|
4
|
+
from .models import RoutingDecision, CostReport
|
|
5
|
+
|
|
6
|
+
__version__ = "2.2.0"
|
|
7
|
+
|
|
8
|
+
__all__ = [
|
|
9
|
+
"A3MRouter",
|
|
10
|
+
"A3MRouterSync",
|
|
11
|
+
"A3MRouterError",
|
|
12
|
+
"RoutingDecision",
|
|
13
|
+
"CostReport",
|
|
14
|
+
# Framework adapters
|
|
15
|
+
"LangChainAdapter",
|
|
16
|
+
"LlamaIndexAdapter",
|
|
17
|
+
"QdrantAdapter",
|
|
18
|
+
"WeaviateAdapter",
|
|
19
|
+
]
|
|
20
|
+
|
|
21
|
+
# Lazy-load adapters
|
|
22
|
+
def __getattr__(name: str):
|
|
23
|
+
if name == "LangChainAdapter":
|
|
24
|
+
from .adapters.langchain import LangChainAdapter
|
|
25
|
+
return LangChainAdapter
|
|
26
|
+
if name == "LlamaIndexAdapter":
|
|
27
|
+
from .adapters.llamaindex import LlamaIndexAdapter
|
|
28
|
+
return LlamaIndexAdapter
|
|
29
|
+
if name == "QdrantAdapter":
|
|
30
|
+
from .adapters.qdrant import QdrantAdapter
|
|
31
|
+
return QdrantAdapter
|
|
32
|
+
if name == "WeaviateAdapter":
|
|
33
|
+
from .adapters.weaviate import WeaviateAdapter
|
|
34
|
+
return WeaviateAdapter
|
|
35
|
+
raise AttributeError(f"module 'a3m' has no attribute '{name}'")
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
"""
|
|
2
|
+
A3M Router framework adapters.
|
|
3
|
+
|
|
4
|
+
Lazy-loaded adapter classes for popular frameworks:
|
|
5
|
+
- LangChain: from a3m.adapters import LangChainAdapter
|
|
6
|
+
- LlamaIndex: from a3m.adapters import LlamaIndexAdapter
|
|
7
|
+
- Qdrant: from a3m.adapters import QdrantAdapter
|
|
8
|
+
- Weaviate: from a3m.adapters import WeaviateAdapter
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from .langchain import LangChainAdapter
|
|
12
|
+
from .llamaindex import LlamaIndexAdapter
|
|
13
|
+
from .qdrant import QdrantAdapter
|
|
14
|
+
from .weaviate import WeaviateAdapter
|
|
15
|
+
|
|
16
|
+
__all__ = [
|
|
17
|
+
"LangChainAdapter",
|
|
18
|
+
"LlamaIndexAdapter",
|
|
19
|
+
"QdrantAdapter",
|
|
20
|
+
"WeaviateAdapter",
|
|
21
|
+
]
|
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
"""
|
|
2
|
+
A3M Router LangChain adapter.
|
|
3
|
+
|
|
4
|
+
Drop-in replacement for ChatOpenAI that routes through A3M Router.
|
|
5
|
+
|
|
6
|
+
Usage:
|
|
7
|
+
from langchain.chat_models import ChatOpenAI
|
|
8
|
+
from a3m.adapters import LangChainAdapter
|
|
9
|
+
|
|
10
|
+
# As replacement for ChatOpenAI
|
|
11
|
+
llm = LangChainAdapter(
|
|
12
|
+
base_url="http://localhost:8787",
|
|
13
|
+
model="auto",
|
|
14
|
+
temperature=0.7,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
from langchain.schema import HumanMessage
|
|
18
|
+
response = llm([HumanMessage(content="What is 2+2?")])
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from __future__ import annotations
|
|
22
|
+
|
|
23
|
+
import logging
|
|
24
|
+
from typing import Any, Dict, List, Optional, Type
|
|
25
|
+
|
|
26
|
+
from pydantic import Field
|
|
27
|
+
|
|
28
|
+
logger = logging.getLogger(__name__)
|
|
29
|
+
|
|
30
|
+
LANCHAIN_AVAILABLE = False
|
|
31
|
+
try:
|
|
32
|
+
from langchain.chat_models import BaseChatModel
|
|
33
|
+
from langchain.schema import (
|
|
34
|
+
BaseMessage,
|
|
35
|
+
ChatResult,
|
|
36
|
+
AIMessage,
|
|
37
|
+
HumanMessage,
|
|
38
|
+
SystemMessage,
|
|
39
|
+
)
|
|
40
|
+
from langchain.callbacks.manager import CallbackManagerForLLMRun
|
|
41
|
+
LANCHAIN_AVAILABLE = True
|
|
42
|
+
except ImportError:
|
|
43
|
+
logger.warning("LangChain not installed. pip install langchain")
|
|
44
|
+
|
|
45
|
+
from a3m.client import A3MRouter, A3MRouterError
|
|
46
|
+
|
|
47
|
+
if LANCHAIN_AVAILABLE:
|
|
48
|
+
class LangChainAdapter(BaseChatModel):
|
|
49
|
+
"""
|
|
50
|
+
LangChain chat model that routes through A3M Router.
|
|
51
|
+
|
|
52
|
+
Drop-in replacement for ChatOpenAI with automatic model selection.
|
|
53
|
+
|
|
54
|
+
Args:
|
|
55
|
+
base_url: A3M Router server URL.
|
|
56
|
+
model: Model name or "auto" for intelligent routing.
|
|
57
|
+
temperature: Sampling temperature.
|
|
58
|
+
max_tokens: Max tokens to generate.
|
|
59
|
+
parallel_ensemble: Number of providers for ensemble calls.
|
|
60
|
+
**kwargs: Additional A3MRouter options.
|
|
61
|
+
"""
|
|
62
|
+
|
|
63
|
+
base_url: str = Field(default="http://localhost:8787")
|
|
64
|
+
model: str = Field(default="auto")
|
|
65
|
+
temperature: float = Field(default=0.7)
|
|
66
|
+
max_tokens: Optional[int] = Field(default=None)
|
|
67
|
+
parallel_ensemble: int = Field(default=1)
|
|
68
|
+
api_key: Optional[str] = Field(default=None)
|
|
69
|
+
|
|
70
|
+
class Config:
|
|
71
|
+
arbitrary_types_allowed = True
|
|
72
|
+
|
|
73
|
+
def _get_router(self) -> A3MRouter:
|
|
74
|
+
"""Get or create A3M Router client."""
|
|
75
|
+
if not hasattr(self, "_router"):
|
|
76
|
+
self._router = A3MRouter(
|
|
77
|
+
base_url=self.base_url,
|
|
78
|
+
api_key=self.api_key,
|
|
79
|
+
default_model=self.model,
|
|
80
|
+
default_temperature=self.temperature,
|
|
81
|
+
default_max_tokens=self.max_tokens,
|
|
82
|
+
parallel_ensemble=self.parallel_ensemble,
|
|
83
|
+
)
|
|
84
|
+
return self._router
|
|
85
|
+
|
|
86
|
+
def _convert_messages(
|
|
87
|
+
self,
|
|
88
|
+
messages: List[BaseMessage],
|
|
89
|
+
) -> List[Dict[str, str]]:
|
|
90
|
+
"""Convert LangChain messages to A3M format."""
|
|
91
|
+
result = []
|
|
92
|
+
for msg in messages:
|
|
93
|
+
if isinstance(msg, HumanMessage):
|
|
94
|
+
result.append({"role": "user", "content": msg.content})
|
|
95
|
+
elif isinstance(msg, AIMessage):
|
|
96
|
+
result.append({"role": "assistant", "content": msg.content})
|
|
97
|
+
elif isinstance(msg, SystemMessage):
|
|
98
|
+
result.append({"role": "system", "content": msg.content})
|
|
99
|
+
else:
|
|
100
|
+
result.append({"role": "user", "content": str(msg.content)})
|
|
101
|
+
return result
|
|
102
|
+
|
|
103
|
+
def _convert_response(
|
|
104
|
+
self,
|
|
105
|
+
response_content: str,
|
|
106
|
+
) -> AIMessage:
|
|
107
|
+
"""Convert A3M response to LangChain message."""
|
|
108
|
+
return AIMessage(content=response_content)
|
|
109
|
+
|
|
110
|
+
@property
|
|
111
|
+
def _llm_type(self) -> str:
|
|
112
|
+
return "a3m-router"
|
|
113
|
+
|
|
114
|
+
def _generate(
|
|
115
|
+
self,
|
|
116
|
+
messages: List[BaseMessage],
|
|
117
|
+
stop: Optional[List[str]] = None,
|
|
118
|
+
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
119
|
+
**kwargs: Any,
|
|
120
|
+
) -> ChatResult:
|
|
121
|
+
"""Generate a chat response."""
|
|
122
|
+
router = self._get_router()
|
|
123
|
+
|
|
124
|
+
a3m_messages = self._convert_messages(messages)
|
|
125
|
+
|
|
126
|
+
try:
|
|
127
|
+
response = router.chat(
|
|
128
|
+
messages=a3m_messages,
|
|
129
|
+
model=self.model,
|
|
130
|
+
temperature=self.temperature,
|
|
131
|
+
max_tokens=self.max_tokens,
|
|
132
|
+
parallel_ensemble=self.parallel_ensemble,
|
|
133
|
+
**kwargs,
|
|
134
|
+
)
|
|
135
|
+
except A3MRouterError as e:
|
|
136
|
+
logger.error(f"A3M Router error: {e}")
|
|
137
|
+
raise
|
|
138
|
+
|
|
139
|
+
return ChatResult(
|
|
140
|
+
generations=[{"message": self._convert_response(response.content), "text": response.content}],
|
|
141
|
+
llm_output={
|
|
142
|
+
"provider": response.provider,
|
|
143
|
+
"model": response.route.model,
|
|
144
|
+
"cost": response.cost,
|
|
145
|
+
"route": str(response.route),
|
|
146
|
+
},
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
async def _agenerate(
|
|
150
|
+
self,
|
|
151
|
+
messages: List[BaseMessage],
|
|
152
|
+
stop: Optional[List[str]] = None,
|
|
153
|
+
run_manager: Optional[CallbackManagerForLLMRun] = None,
|
|
154
|
+
**kwargs: Any,
|
|
155
|
+
) -> ChatResult:
|
|
156
|
+
"""Async generate a chat response."""
|
|
157
|
+
router = self._get_router()
|
|
158
|
+
a3m_messages = self._convert_messages(messages)
|
|
159
|
+
|
|
160
|
+
try:
|
|
161
|
+
response = await router.achat(
|
|
162
|
+
messages=a3m_messages,
|
|
163
|
+
model=self.model,
|
|
164
|
+
temperature=self.temperature,
|
|
165
|
+
max_tokens=self.max_tokens,
|
|
166
|
+
parallel_ensemble=self.parallel_ensemble,
|
|
167
|
+
**kwargs,
|
|
168
|
+
)
|
|
169
|
+
except A3MRouterError as e:
|
|
170
|
+
logger.error(f"A3M Router error: {e}")
|
|
171
|
+
raise
|
|
172
|
+
|
|
173
|
+
return ChatResult(
|
|
174
|
+
generations=[{"message": self._convert_response(response.content), "text": response.content}],
|
|
175
|
+
llm_output={
|
|
176
|
+
"provider": response.provider,
|
|
177
|
+
"model": response.route.model,
|
|
178
|
+
"cost": response.cost,
|
|
179
|
+
"route": str(response.route),
|
|
180
|
+
},
|
|
181
|
+
)
|
|
182
|
+
|
|
183
|
+
else:
|
|
184
|
+
# Stub class when LangChain is not installed
|
|
185
|
+
class LangChainAdapter:
|
|
186
|
+
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
187
|
+
raise ImportError(
|
|
188
|
+
"LangChain is not installed. "
|
|
189
|
+
"Install with: pip install langchain"
|
|
190
|
+
)
|
|
@@ -0,0 +1,249 @@
|
|
|
1
|
+
"""
|
|
2
|
+
A3M Router LlamaIndex adapter.
|
|
3
|
+
|
|
4
|
+
Drop-in LLM for LlamaIndex that routes through A3M Router.
|
|
5
|
+
|
|
6
|
+
Usage:
|
|
7
|
+
from llama_index import VectorStoreIndex, SimpleWebPageReader
|
|
8
|
+
from a3m.adapters import LlamaIndexAdapter
|
|
9
|
+
|
|
10
|
+
llm = LlamaIndexAdapter(
|
|
11
|
+
base_url="http://localhost:8787",
|
|
12
|
+
model="auto",
|
|
13
|
+
temperature=0.7,
|
|
14
|
+
)
|
|
15
|
+
|
|
16
|
+
index = VectorStoreIndex.from_documents(
|
|
17
|
+
documents,
|
|
18
|
+
llm=llm, # Use A3M Router as the LLM
|
|
19
|
+
)
|
|
20
|
+
"""
|
|
21
|
+
|
|
22
|
+
from __future__ import annotations
|
|
23
|
+
|
|
24
|
+
import logging
|
|
25
|
+
from typing import Any, Awaitable, List, Optional
|
|
26
|
+
|
|
27
|
+
logger = logging.getLogger(__name__)
|
|
28
|
+
|
|
29
|
+
LLAMAINDEX_AVAILABLE = False
|
|
30
|
+
try:
|
|
31
|
+
from llama_index.llms import BaseLLM
|
|
32
|
+
from llama_index.llms.custom import CustomLLM
|
|
33
|
+
from llama_index.types import ModelType
|
|
34
|
+
from llama_index.output_parsers.base import BaseOutputParser
|
|
35
|
+
LLAMAINDEX_AVAILABLE = True
|
|
36
|
+
except ImportError:
|
|
37
|
+
logger.warning("LlamaIndex not installed. pip install llama-index")
|
|
38
|
+
|
|
39
|
+
from a3m.client import A3MRouter, A3MRouterError
|
|
40
|
+
|
|
41
|
+
if LLAMAINDEX_AVAILABLE:
|
|
42
|
+
class LlamaIndexAdapter(BaseLLM):
|
|
43
|
+
"""
|
|
44
|
+
LlamaIndex LLM that routes through A3M Router.
|
|
45
|
+
|
|
46
|
+
Drop-in replacement for OpenAI/GPT LLMs in LlamaIndex pipelines.
|
|
47
|
+
|
|
48
|
+
Args:
|
|
49
|
+
base_url: A3M Router server URL.
|
|
50
|
+
model: Model name or "auto" for intelligent routing.
|
|
51
|
+
temperature: Sampling temperature.
|
|
52
|
+
max_tokens: Max tokens to generate.
|
|
53
|
+
parallel_ensemble: Number of providers for ensemble calls.
|
|
54
|
+
"""
|
|
55
|
+
|
|
56
|
+
base_url: str = "http://localhost:8787"
|
|
57
|
+
model: str = "auto"
|
|
58
|
+
temperature: float = 0.7
|
|
59
|
+
max_tokens: Optional[int] = 4096
|
|
60
|
+
parallel_ensemble: int = 1
|
|
61
|
+
api_key: Optional[str] = None
|
|
62
|
+
|
|
63
|
+
def __init__(self, **kwargs: Any) -> None:
|
|
64
|
+
super().__init__(**kwargs)
|
|
65
|
+
self._router: Optional[A3MRouter] = None
|
|
66
|
+
|
|
67
|
+
def _get_router(self) -> A3MRouter:
|
|
68
|
+
"""Get or create A3M Router client."""
|
|
69
|
+
if self._router is None:
|
|
70
|
+
self._router = A3MRouter(
|
|
71
|
+
base_url=self.base_url,
|
|
72
|
+
api_key=self.api_key,
|
|
73
|
+
default_model=self.model,
|
|
74
|
+
default_temperature=self.temperature,
|
|
75
|
+
default_max_tokens=self.max_tokens,
|
|
76
|
+
parallel_ensemble=self.parallel_ensemble,
|
|
77
|
+
)
|
|
78
|
+
return self._router
|
|
79
|
+
|
|
80
|
+
@property
|
|
81
|
+
def model_type(self) -> ModelType:
|
|
82
|
+
return ModelType.LLM
|
|
83
|
+
|
|
84
|
+
@property
|
|
85
|
+
def class_name(self) -> str:
|
|
86
|
+
return "A3MRouter"
|
|
87
|
+
|
|
88
|
+
def complete(
|
|
89
|
+
self,
|
|
90
|
+
prompt: str,
|
|
91
|
+
formatted: bool = False,
|
|
92
|
+
**kwargs: Any,
|
|
93
|
+
) -> Any:
|
|
94
|
+
"""
|
|
95
|
+
Synchronous completion.
|
|
96
|
+
|
|
97
|
+
LlamaIndex calls this for text completion.
|
|
98
|
+
"""
|
|
99
|
+
router = self._get_router()
|
|
100
|
+
|
|
101
|
+
try:
|
|
102
|
+
response = router.chat(
|
|
103
|
+
messages=[{"role": "user", "content": prompt}],
|
|
104
|
+
model=self.model,
|
|
105
|
+
temperature=self.temperature,
|
|
106
|
+
max_tokens=self.max_tokens,
|
|
107
|
+
parallel_ensemble=self.parallel_ensemble,
|
|
108
|
+
**kwargs,
|
|
109
|
+
)
|
|
110
|
+
return response.content
|
|
111
|
+
except A3MRouterError as e:
|
|
112
|
+
logger.error(f"A3M Router error: {e}")
|
|
113
|
+
raise
|
|
114
|
+
|
|
115
|
+
async def acomplete(
|
|
116
|
+
self,
|
|
117
|
+
prompt: str,
|
|
118
|
+
formatted: bool = False,
|
|
119
|
+
**kwargs: Any,
|
|
120
|
+
) -> Any:
|
|
121
|
+
"""Async completion."""
|
|
122
|
+
router = self._get_router()
|
|
123
|
+
|
|
124
|
+
try:
|
|
125
|
+
response = await router.achat(
|
|
126
|
+
messages=[{"role": "user", "content": prompt}],
|
|
127
|
+
model=self.model,
|
|
128
|
+
temperature=self.temperature,
|
|
129
|
+
max_tokens=self.max_tokens,
|
|
130
|
+
parallel_ensemble=self.parallel_ensemble,
|
|
131
|
+
**kwargs,
|
|
132
|
+
)
|
|
133
|
+
return response.content
|
|
134
|
+
except A3MRouterError as e:
|
|
135
|
+
logger.error(f"A3M Router error: {e}")
|
|
136
|
+
raise
|
|
137
|
+
|
|
138
|
+
def stream_complete(
|
|
139
|
+
self,
|
|
140
|
+
prompt: str,
|
|
141
|
+
**kwargs: Any,
|
|
142
|
+
) -> Any:
|
|
143
|
+
"""
|
|
144
|
+
Streaming completion.
|
|
145
|
+
|
|
146
|
+
Returns a generator that yields response chunks.
|
|
147
|
+
"""
|
|
148
|
+
router = self._get_router()
|
|
149
|
+
|
|
150
|
+
try:
|
|
151
|
+
chunks = router.stream_chat(
|
|
152
|
+
messages=[{"role": "user", "content": prompt}],
|
|
153
|
+
model=self.model,
|
|
154
|
+
temperature=self.temperature,
|
|
155
|
+
max_tokens=self.max_tokens,
|
|
156
|
+
**kwargs,
|
|
157
|
+
)
|
|
158
|
+
for chunk in chunks:
|
|
159
|
+
yield chunk.content
|
|
160
|
+
except A3MRouterError as e:
|
|
161
|
+
logger.error(f"A3M Router error: {e}")
|
|
162
|
+
raise
|
|
163
|
+
|
|
164
|
+
async def astream_complete(
|
|
165
|
+
self,
|
|
166
|
+
prompt: str,
|
|
167
|
+
**kwargs: Any,
|
|
168
|
+
) -> Awaitable[Any]:
|
|
169
|
+
"""Async streaming completion."""
|
|
170
|
+
router = self._get_router()
|
|
171
|
+
|
|
172
|
+
async def gen():
|
|
173
|
+
async for chunk in router.astream_chat(
|
|
174
|
+
messages=[{"role": "user", "content": prompt}],
|
|
175
|
+
model=self.model,
|
|
176
|
+
temperature=self.temperature,
|
|
177
|
+
max_tokens=self.max_tokens,
|
|
178
|
+
**kwargs,
|
|
179
|
+
):
|
|
180
|
+
yield chunk.content
|
|
181
|
+
|
|
182
|
+
return gen()
|
|
183
|
+
|
|
184
|
+
def chat(self, messages: List[Any], **kwargs: Any) -> Any:
|
|
185
|
+
"""
|
|
186
|
+
Synchronous chat.
|
|
187
|
+
|
|
188
|
+
Converts messages to a single prompt.
|
|
189
|
+
"""
|
|
190
|
+
router = self._get_router()
|
|
191
|
+
|
|
192
|
+
# Convert message objects to content strings
|
|
193
|
+
if hasattr(messages[0], "content"):
|
|
194
|
+
content = "\n".join([getattr(m, "content", str(m)) for m in messages])
|
|
195
|
+
else:
|
|
196
|
+
content = str(messages[0])
|
|
197
|
+
|
|
198
|
+
try:
|
|
199
|
+
response = router.chat(
|
|
200
|
+
messages=[{"role": "user", "content": content}],
|
|
201
|
+
model=self.model,
|
|
202
|
+
temperature=self.temperature,
|
|
203
|
+
max_tokens=self.max_tokens,
|
|
204
|
+
**kwargs,
|
|
205
|
+
)
|
|
206
|
+
# Return a simple object that LlamaIndex expects
|
|
207
|
+
return ChatMessage(content=response.content)
|
|
208
|
+
except A3MRouterError as e:
|
|
209
|
+
logger.error(f"A3M Router error: {e}")
|
|
210
|
+
raise
|
|
211
|
+
|
|
212
|
+
async def achat(self, messages: List[Any], **kwargs: Any) -> Any:
|
|
213
|
+
"""Async chat."""
|
|
214
|
+
router = self._get_router()
|
|
215
|
+
|
|
216
|
+
if hasattr(messages[0], "content"):
|
|
217
|
+
content = "\n".join([getattr(m, "content", str(m)) for m in messages])
|
|
218
|
+
else:
|
|
219
|
+
content = str(messages[0])
|
|
220
|
+
|
|
221
|
+
try:
|
|
222
|
+
response = await router.achat(
|
|
223
|
+
messages=[{"role": "user", "content": content}],
|
|
224
|
+
model=self.model,
|
|
225
|
+
temperature=self.temperature,
|
|
226
|
+
max_tokens=self.max_tokens,
|
|
227
|
+
**kwargs,
|
|
228
|
+
)
|
|
229
|
+
return ChatMessage(content=response.content)
|
|
230
|
+
except A3MRouterError as e:
|
|
231
|
+
logger.error(f"A3M Router error: {e}")
|
|
232
|
+
raise
|
|
233
|
+
|
|
234
|
+
else:
|
|
235
|
+
# Stub class when LlamaIndex is not installed
|
|
236
|
+
class LlamaIndexAdapter:
|
|
237
|
+
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
238
|
+
raise ImportError(
|
|
239
|
+
"LlamaIndex is not installed. "
|
|
240
|
+
"Install with: pip install llama-index"
|
|
241
|
+
)
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
# Simple message class for LlamaIndex compatibility
|
|
245
|
+
class ChatMessage:
|
|
246
|
+
"""Simple chat message for LlamaIndex compatibility."""
|
|
247
|
+
def __init__(self, content: str) -> None:
|
|
248
|
+
self.content = content
|
|
249
|
+
self.raw = {}
|