ai-microcore 4.0.0.dev23__tar.gz → 4.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.
Files changed (45) hide show
  1. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/PKG-INFO +27 -5
  2. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/README.md +23 -2
  3. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/__init__.py +1 -1
  4. ai_microcore-4.2.0/microcore/_mcp.py +381 -0
  5. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/embedding_db/qdrant.py +40 -6
  6. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/logging.py +1 -1
  7. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/mcp.py +25 -1
  8. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/ui.py +1 -1
  9. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/pyproject.toml +3 -2
  10. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/LICENSE +0 -0
  11. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/_env.py +0 -0
  12. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/_llm_functions.py +0 -0
  13. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/_prepare_llm_args.py +0 -0
  14. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/ai_func/__init__.py +0 -0
  15. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/ai_func/ai-func.json.j2 +0 -0
  16. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/ai_func/ai-func.pythonic.j2 +0 -0
  17. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/ai_modules.py +0 -0
  18. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/configuration.py +0 -0
  19. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/embedding_db/__init__.py +0 -0
  20. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/embedding_db/chromadb.py +0 -0
  21. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/file_storage.py +0 -0
  22. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/interactive_setup.py +0 -0
  23. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/json_parsing.py +0 -0
  24. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/__init__.py +0 -0
  25. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/_openai_llm_v0.py +0 -0
  26. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/_openai_llm_v1.py +0 -0
  27. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/anthropic.py +0 -0
  28. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/google_genai.py +0 -0
  29. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/google_vertex_ai.py +0 -0
  30. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/local_llm.py +0 -0
  31. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/local_transformers.py +0 -0
  32. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/openai_llm.py +0 -0
  33. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/llm/shared.py +0 -0
  34. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/message_types.py +0 -0
  35. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/metrics.py +0 -0
  36. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/python.py +0 -0
  37. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/templating/__init__.py +0 -0
  38. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/templating/jinja2.py +0 -0
  39. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/text2speech/elevenlabs.py +0 -0
  40. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/tokenizing.py +0 -0
  41. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/types.py +0 -0
  42. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/utils.py +0 -0
  43. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/wrappers/__init__.py +0 -0
  44. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/wrappers/llm_response_wrapper.py +0 -0
  45. {ai_microcore-4.0.0.dev23 → ai_microcore-4.2.0}/microcore/wrappers/prompt_wrapper.py +0 -0
@@ -1,10 +1,10 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ai-microcore
3
- Version: 4.0.0.dev23
3
+ Version: 4.2.0
4
4
  Summary: # Minimalistic Foundation for AI Applications
5
5
  Keywords: llm,large language models,ai,similarity search,ai search,gpt,openai,framework,adapter
6
- Author-email: Vitalii Stepanenko <mail@vitalii.in>
7
- Maintainer-email: Vitalii Stepanenko <mail@vitalii.in>
6
+ Author-email: Vitalii Stepanenko <mail@vitaliy.in>
7
+ Maintainer-email: Vitalii Stepanenko <mail@vitaliy.in>
8
8
  Requires-Python: >=3.10
9
9
  Description-Content-Type: text/markdown
10
10
  Classifier: Programming Language :: Python :: 3
@@ -26,13 +26,14 @@ Requires-Dist: tiktoken>=0.7.0,<1.0
26
26
  Requires-Dist: mcp~=1.9.2
27
27
  Requires-Dist: fastmcp~=2.8.0
28
28
  Requires-Dist: docstring_parser~=0.16.0
29
+ Requires-Dist: httpx~=0.28.1
29
30
  Project-URL: Source Code, https://github.com/Nayjest/ai-microcore
30
31
 
31
32
  <p align="right">
32
33
  <a href="https://github.com/Nayjest/ai-microcore/releases" target="_blank"><img src="https://img.shields.io/github/v/release/Nayjest/ai-microcore.svg" alt="Release Notes"></a>
33
34
  <a href="https://app.codacy.com/gh/Nayjest/ai-microcore/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_grade" target="_blank"><img src="https://app.codacy.com/project/badge/Grade/441d03416bc048828c649129530dcbc3" alt="Code Quality"></a>
34
35
  <a href="https://github.com/Nayjest/ai-microcore/actions/workflows/pylint.yml" target="_blank"><img src="https://github.com/Nayjest/ai-microcore/actions/workflows/pylint.yml/badge.svg" alt="Pylint"></a>
35
- <img src="https://github.com/Nayjest/ai-microcore/blob/main/coverage.svg" alt="Code Coverage">
36
+ <img src="https://raw.githubusercontent.com/Nayjest/ai-microcore/main/coverage.svg" alt="Code Coverage">
36
37
  <a href="https://github.com/Nayjest/ai-microcore/actions/workflows/tests.yml" target="_blank"><img src="https://github.com/Nayjest/ai-microcore/actions/workflows/tests.yml/badge.svg" alt="Tests"></a>
37
38
  <a href="https://github.com/Nayjest/ai-microcore/blob/main/LICENSE" target="_blank"><img src="https://img.shields.io/static/v1?label=license&message=MIT&color=d08aff" alt="License"></a>
38
39
  </p>
@@ -142,8 +143,10 @@ See [transformers installation](https://huggingface.co/docs/transformers/install
142
143
  ### Vector Databases
143
144
 
144
145
  Vector database functions are available via `microcore.texts`.
146
+
147
+ #### ChromaDB
145
148
  Default vector database is [Chroma](https://www.trychroma.com/).
146
- In order to use vector database functions, you need to install the `chromadb` package:
149
+ In order to use vector database functions with ChromaDB, you need to install the `chromadb` package:
147
150
  ```bash
148
151
  pip install chromadb
149
152
  ```
@@ -157,6 +160,25 @@ configure(
157
160
  EMBEDDING_DB_PORT = 8000,
158
161
  )
159
162
  ```
163
+ #### Qdrant
164
+ In order to use vector database functions with Qdrant, you need to install the `qdrant-client` package:
165
+ ```bash
166
+ pip install qdrant-client
167
+ ```
168
+ Configuration example
169
+ ```python
170
+ from microcore import configure, EmbeddingDbType
171
+ from sentence_transformers import SentenceTransformer
172
+
173
+ configure(
174
+ EMBEDDING_DB_TYPE=EmbeddingDbType.QDRANT,
175
+ EMBEDDING_DB_HOST="localhost",
176
+ EMBEDDING_DB_PORT="6333",
177
+ EMBEDDING_DB_SIZE=384, # dimentions quantity in used SentenceTransformer model
178
+ EMBEDDING_DB_FUNCTION=SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2"),
179
+ )
180
+ ```
181
+
160
182
 
161
183
  ## 🌟 Core Functions
162
184
 
@@ -2,7 +2,7 @@
2
2
  <a href="https://github.com/Nayjest/ai-microcore/releases" target="_blank"><img src="https://img.shields.io/github/v/release/Nayjest/ai-microcore.svg" alt="Release Notes"></a>
3
3
  <a href="https://app.codacy.com/gh/Nayjest/ai-microcore/dashboard?utm_source=gh&utm_medium=referral&utm_content=&utm_campaign=Badge_grade" target="_blank"><img src="https://app.codacy.com/project/badge/Grade/441d03416bc048828c649129530dcbc3" alt="Code Quality"></a>
4
4
  <a href="https://github.com/Nayjest/ai-microcore/actions/workflows/pylint.yml" target="_blank"><img src="https://github.com/Nayjest/ai-microcore/actions/workflows/pylint.yml/badge.svg" alt="Pylint"></a>
5
- <img src="https://github.com/Nayjest/ai-microcore/blob/main/coverage.svg" alt="Code Coverage">
5
+ <img src="https://raw.githubusercontent.com/Nayjest/ai-microcore/main/coverage.svg" alt="Code Coverage">
6
6
  <a href="https://github.com/Nayjest/ai-microcore/actions/workflows/tests.yml" target="_blank"><img src="https://github.com/Nayjest/ai-microcore/actions/workflows/tests.yml/badge.svg" alt="Tests"></a>
7
7
  <a href="https://github.com/Nayjest/ai-microcore/blob/main/LICENSE" target="_blank"><img src="https://img.shields.io/static/v1?label=license&message=MIT&color=d08aff" alt="License"></a>
8
8
  </p>
@@ -112,8 +112,10 @@ See [transformers installation](https://huggingface.co/docs/transformers/install
112
112
  ### Vector Databases
113
113
 
114
114
  Vector database functions are available via `microcore.texts`.
115
+
116
+ #### ChromaDB
115
117
  Default vector database is [Chroma](https://www.trychroma.com/).
116
- In order to use vector database functions, you need to install the `chromadb` package:
118
+ In order to use vector database functions with ChromaDB, you need to install the `chromadb` package:
117
119
  ```bash
118
120
  pip install chromadb
119
121
  ```
@@ -127,6 +129,25 @@ configure(
127
129
  EMBEDDING_DB_PORT = 8000,
128
130
  )
129
131
  ```
132
+ #### Qdrant
133
+ In order to use vector database functions with Qdrant, you need to install the `qdrant-client` package:
134
+ ```bash
135
+ pip install qdrant-client
136
+ ```
137
+ Configuration example
138
+ ```python
139
+ from microcore import configure, EmbeddingDbType
140
+ from sentence_transformers import SentenceTransformer
141
+
142
+ configure(
143
+ EMBEDDING_DB_TYPE=EmbeddingDbType.QDRANT,
144
+ EMBEDDING_DB_HOST="localhost",
145
+ EMBEDDING_DB_PORT="6333",
146
+ EMBEDDING_DB_SIZE=384, # dimentions quantity in used SentenceTransformer model
147
+ EMBEDDING_DB_FUNCTION=SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2"),
148
+ )
149
+ ```
150
+
130
151
 
131
152
  ## 🌟 Core Functions
132
153
 
@@ -186,4 +186,4 @@ __all__ = [
186
186
  # "wrappers",
187
187
  ]
188
188
 
189
- __version__ = "4.0.0-dev23"
189
+ __version__ = "4.2.0"
@@ -0,0 +1,381 @@
1
+ import asyncio
2
+ import logging
3
+ import datetime
4
+ from typing import Optional, Dict, Any
5
+ from dataclasses import dataclass, field
6
+ from enum import Enum
7
+
8
+ import requests
9
+ from fastmcp import Client
10
+ from fastmcp.client.progress import ProgressHandler
11
+ import mcp.types
12
+
13
+ from .utils import ExtendedString, ConvertableToMessage
14
+ from .ai_func import AiFuncSyntax
15
+ from . import ui
16
+ from .types import BadAIAnswer, BadAIJsonAnswer
17
+ from .wrappers.llm_response_wrapper import LLMResponse
18
+ from ._env import env
19
+ from .file_storage import storage
20
+
21
+
22
+ class WrongMcpUsage(BadAIAnswer):
23
+ ...
24
+
25
+
26
+ class ToolsCache:
27
+ FILE = "mcp/tools.json"
28
+
29
+ @staticmethod
30
+ def read(mcp_url: str) -> Optional["Tools"]:
31
+ logging.info(f"Checking MCP tools cache for {ui.green(mcp_url)}...")
32
+ tools = storage.read_json(ToolsCache.FILE, default={}).get(mcp_url, None)
33
+ if tools is not None:
34
+ return Tools.from_list([Tool(**raw) for raw in tools.values()])
35
+ return None
36
+
37
+ @staticmethod
38
+ def write(mcp_url: str, tools: "Tools"):
39
+ logging.info(f"Storing MCP tools cache for {ui.green(mcp_url)}...")
40
+ cached_tools = storage.read_json(ToolsCache.FILE, default={})
41
+ cached_tools[mcp_url] = tools
42
+ storage.write_json(ToolsCache.FILE, cached_tools)
43
+
44
+ @staticmethod
45
+ def clear():
46
+ logging.info("Clearing MCP tools cache...")
47
+ storage.delete(ToolsCache.FILE)
48
+
49
+
50
+ class MCPAnswer(ExtendedString, ConvertableToMessage):
51
+ ...
52
+
53
+
54
+ class McpTransport(str, Enum):
55
+ SSE: str = "sse"
56
+ STREAMABLE_HTTP: str = "streamable_http"
57
+ WS: str = "ws"
58
+ STDIO: str = "stdio"
59
+
60
+ def __str__(self):
61
+ return self.value
62
+
63
+
64
+ @dataclass
65
+ class MCPConnection:
66
+ url: str = None
67
+ transport: McpTransport = field(default=McpTransport.STREAMABLE_HTTP)
68
+ headers: Optional[Dict[str, str]] = field(default_factory=dict)
69
+ tools: Optional["Tools"] = field(default=None, init=False)
70
+ _client: Client | None = field(default=None, init=False)
71
+
72
+ @staticmethod
73
+ async def init(
74
+ url: str,
75
+ transport: McpTransport,
76
+ headers: Optional[Dict[str, str]] = None,
77
+ fetch_tools: bool = True,
78
+ use_cache: bool = True,
79
+ connect_timeout: float = 10,
80
+ ) -> "MCPConnection":
81
+ con: MCPConnection = MCPConnection(url=url, transport=transport, headers=headers)
82
+ con._client = Client(url, timeout=connect_timeout, headers=headers) # pylint: disable=W0212
83
+ await con._client.__aenter__() # pylint: disable=E1101,W0212,C2801
84
+ if fetch_tools:
85
+ await con.fetch_tools(use_cache=use_cache)
86
+ return con
87
+
88
+ async def close(self):
89
+ if self._client:
90
+ await self._client.__aexit__(None, None, None)
91
+ self._client = None
92
+ else:
93
+ logging.error(f"Trying to close MCP connection that is not opened ({self.url})")
94
+
95
+ async def fetch_tools(self, use_cache: bool = True) -> "Tools":
96
+ if self.tools is not None:
97
+ return self.tools
98
+
99
+ if use_cache and (cached_tools := ToolsCache.read(self.url)):
100
+ logging.info("Using MCP tools from cache for %s", ui.green(self.url))
101
+ self.tools = cached_tools
102
+ return self.tools
103
+
104
+ logging.info("Fetching tools from MCP %s", ui.green(self.url))
105
+ mcp_tools = await self._client.list_tools()
106
+ self.tools = Tools.from_list([Tool.from_mcp(tool) for tool in mcp_tools])
107
+ if use_cache:
108
+ self.update_tools_cache()
109
+ return self.tools
110
+
111
+ def update_tools_cache(self):
112
+ if self.tools is None:
113
+ raise RuntimeError("Tools are not fetched yet. Call fetch_tools() first.")
114
+ ToolsCache.write(self.url, self.tools)
115
+
116
+ async def call(
117
+ self,
118
+ name: str,
119
+ timeout: datetime.timedelta | float | int | None = None,
120
+ progress_handler: ProgressHandler | None = None,
121
+ **kwargs
122
+ ):
123
+ assert env().config.AI_SYNTAX_FUNCTION_NAME_FIELD not in kwargs
124
+ params = dict(kwargs)
125
+ params[env().config.AI_SYNTAX_FUNCTION_NAME_FIELD] = name
126
+ return await self.exec(params, timeout=timeout, progress_handler=progress_handler)
127
+
128
+ async def exec(
129
+ self,
130
+ params: dict | LLMResponse,
131
+ timeout: datetime.timedelta | float | int | None = None,
132
+ progress_handler: ProgressHandler | None = None,
133
+ ):
134
+ if isinstance(params, LLMResponse):
135
+ try:
136
+ params = params.parse_json(
137
+ raise_errors=True,
138
+ required_fields=[env().config.AI_SYNTAX_FUNCTION_NAME_FIELD],
139
+ )
140
+ except BadAIJsonAnswer as e:
141
+ raise WrongMcpUsage(str(e)) from e
142
+ params = dict(params)
143
+ name = params.pop(env().config.AI_SYNTAX_FUNCTION_NAME_FIELD)
144
+ if not name:
145
+ raise WrongMcpUsage(
146
+ f"Tool name should be passed in {env().config.AI_SYNTAX_FUNCTION_NAME_FIELD} field"
147
+ )
148
+ logging.info(f"Calling MCP tool {ui.green(name)} with {params}...")
149
+ content = await self._client.call_tool(
150
+ name=name,
151
+ arguments=params,
152
+ timeout=timeout,
153
+ progress_handler=progress_handler
154
+ )
155
+ if content and len(content) == 1 and content[0].type == "text":
156
+ return MCPAnswer(content[0].text, dict(response=content))
157
+ return content
158
+
159
+
160
+ @dataclass
161
+ class Tool:
162
+ @dataclass
163
+ class Arg:
164
+ name: str = field()
165
+ description: str = field(default="")
166
+ type: str = field(default="string")
167
+ required: bool = field(default=...)
168
+ default: any = field(default=...)
169
+
170
+ def __post_init__(self):
171
+ if self.required is ...:
172
+ self.required = self.default is ...
173
+ if self.default is ...:
174
+ self.default = None
175
+
176
+ name: str = field()
177
+ description: str = field(default="")
178
+ args: dict[str, Arg | dict] = field(default_factory=dict)
179
+
180
+ def __post_init__(self):
181
+ for key in self.args.keys(): # pylint: disable=C0201, C0206
182
+ if isinstance(self.args[key], dict):
183
+ self.args[key] = Tool.Arg(**self.args[key])
184
+
185
+ @staticmethod
186
+ def from_mcp(tool: mcp.types.Tool) -> "Tool":
187
+ t = Tool(
188
+ name=tool.name,
189
+ description=tool.description,
190
+ )
191
+ for param_name, data in tool.inputSchema.get("properties", {}).items():
192
+ param = Tool.Arg(
193
+ name=param_name,
194
+ description=data.get("title", ""),
195
+ type=data.get("type", "string"),
196
+ required=param_name in tool.inputSchema.get("required", []),
197
+ default="",
198
+ )
199
+ t.args[param_name] = param
200
+ return t
201
+
202
+ def describe(self, syntax: AiFuncSyntax = None) -> str:
203
+ syntax = syntax or AiFuncSyntax.DEFAULT
204
+ tpl_file = f"ai-func.{syntax}.j2" if syntax in AiFuncSyntax else syntax
205
+ metadata = self._get_metadata()
206
+ return env().tpl_function(tpl_file, **metadata)
207
+
208
+ def _get_metadata(self):
209
+ return dict(
210
+ name=self.name,
211
+ description=self.description,
212
+ args={
213
+ arg.name: dict(
214
+ default="NOT_SET" if arg.required else arg.default,
215
+ type=arg.type,
216
+ comment=arg.description, # @todo rename it to description
217
+ )
218
+ for arg in self.args.values()
219
+ }
220
+ )
221
+
222
+ def __str__(self):
223
+ return self.describe(syntax=AiFuncSyntax.DEFAULT)
224
+
225
+
226
+ class Tools(dict[str, Tool]):
227
+ def __str__(self):
228
+ return "\n".join([str(tool) for tool in self.values()])
229
+
230
+ @staticmethod
231
+ def from_list(tools: list[Tool]) -> "Tools":
232
+ return Tools({tool.name: tool for tool in tools})
233
+
234
+
235
+ @dataclass
236
+ class MCPServer:
237
+ url: str
238
+ name: str = field(default="")
239
+ auth_token: Optional[str] = field(default=None)
240
+ headers: Optional[Dict[str, str]] = field(default=None)
241
+ tools: Tools = field(default_factory=Tools)
242
+ transport: McpTransport = field(default=None)
243
+
244
+ @staticmethod
245
+ def _try_sse_or_streamable_http(url, headers=None) -> tuple[McpTransport, str]:
246
+ if url.endswith("/"):
247
+ test_sse_url = f"{url}sse"
248
+ else:
249
+ test_sse_url = f"{url}/sse"
250
+ try:
251
+ request_headers = headers or {}
252
+ response = requests.request(
253
+ method="HEAD",
254
+ url=test_sse_url,
255
+ headers=request_headers,
256
+ timeout=5
257
+ )
258
+ if response.status_code == 200:
259
+ return McpTransport.SSE, test_sse_url
260
+ except requests.RequestException:
261
+ # not a SSE endpoint or not reachable
262
+ pass
263
+ return McpTransport.STREAMABLE_HTTP, f"{url}/mcp"
264
+
265
+ @staticmethod
266
+ def _guess_transport_type_by_url(url: str) -> Optional[McpTransport]:
267
+ if url.startswith("ws://") or url.startswith("wss://"):
268
+ return McpTransport.WS
269
+ if url.endswith("/mcp"):
270
+ return McpTransport.STREAMABLE_HTTP
271
+ if url.endswith("/sse"):
272
+ return McpTransport.SSE
273
+ return None
274
+
275
+ @staticmethod
276
+ def name_from_url(url: str) -> str:
277
+ """Domain name from URL."""
278
+ return url.split("//")[-1].split("/")[0]
279
+
280
+ def __post_init__(self):
281
+ if not self.name:
282
+ self.name = MCPServer.name_from_url(self.url)
283
+ if not self.transport:
284
+ self.transport = self._guess_transport_type_by_url(self.url)
285
+
286
+ # Prepare headers with authentication if provided
287
+ if self.headers is None:
288
+ self.headers = {}
289
+
290
+ if self.auth_token:
291
+ self.headers["Authorization"] = f"Bearer {self.auth_token}"
292
+
293
+ async def connect(
294
+ self,
295
+ fetch_tools: bool = True,
296
+ use_cache: bool = True,
297
+ connect_timeout: float = 10,
298
+ ) -> MCPConnection:
299
+ if self.transport:
300
+ transport, url = self.transport, self.url
301
+ else:
302
+ transport, url = self._try_sse_or_streamable_http(self.url, self.headers)
303
+
304
+ return await MCPConnection.init(
305
+ url=url,
306
+ transport=transport,
307
+ auth_token=self.auth_token,
308
+ headers=self.headers,
309
+ fetch_tools=fetch_tools,
310
+ use_cache=use_cache,
311
+ connect_timeout=connect_timeout,
312
+ )
313
+
314
+ def get_tools_cache(self) -> Tools | None:
315
+ return ToolsCache.read(self.url)
316
+
317
+
318
+ class MCPRegistry(dict[str, MCPServer]):
319
+ def __init__(self, server_configs: list[dict | str]):
320
+ super().__init__()
321
+ for server_config in server_configs:
322
+ if isinstance(server_config, str):
323
+ server_config = {
324
+ "name": MCPServer.name_from_url(server_config),
325
+ "url": server_config
326
+ }
327
+ self[server_config["name"]] = MCPServer(**server_config)
328
+
329
+ def get(self, server_name: str) -> MCPServer:
330
+ if server_name not in self:
331
+ raise ValueError(f"MCP server '{server_name}' not found in registry")
332
+ return self[server_name]
333
+
334
+ async def precache_tools(
335
+ self,
336
+ raise_errors: bool = False,
337
+ connect_timeout: int = 10,
338
+ ):
339
+ async def precache_server_tools(server_name):
340
+ conn = None
341
+ try:
342
+ conn = await self.get(server_name).connect(
343
+ fetch_tools=True,
344
+ use_cache=False,
345
+ connect_timeout=connect_timeout,
346
+ )
347
+ conn.update_tools_cache()
348
+ except Exception as e: # pylint: disable=W0718
349
+ logging.error("Failed to precache tools for MCP server %s: %s", server_name, e)
350
+ if raise_errors:
351
+ raise
352
+ finally:
353
+ if conn is not None:
354
+ await conn.close()
355
+
356
+ await asyncio.gather(*[precache_server_tools(srv) for srv in self.keys()])
357
+
358
+ async def connect_to(
359
+ self,
360
+ server_name: str,
361
+ fetch_tools: bool = True,
362
+ use_cache: bool = True,
363
+ ) -> MCPConnection:
364
+ mcp_server = self.get(server_name)
365
+ return await mcp_server.connect(fetch_tools=fetch_tools, use_cache=use_cache)
366
+
367
+
368
+ def server(name: str) -> MCPServer:
369
+ """
370
+ Returns MCP server by name from the registry.
371
+
372
+ Args:
373
+ name (str): The name of the MCP server.
374
+
375
+ Returns:
376
+ MCPServer: The MCP server instance.
377
+
378
+ Raises:
379
+ ValueError: If the server with the given name is not found in the registry.
380
+ """
381
+ return env().mcp_registry.get(name)
@@ -67,21 +67,55 @@ class QdrantEmbeddingDB(AbstractEmbeddingDB):
67
67
  )
68
68
 
69
69
  @classmethod
70
- def _convert_where(cls, where: dict | None, kwargs=None) -> Filter | None:
70
+ def _convert_where( # pylint: disable=too-many-branches
71
+ cls,
72
+ where: dict | None,
73
+ kwargs=None
74
+ ) -> Filter | None:
75
+ where_doc = kwargs and kwargs.get("where_document", {}).get("$contains", None)
76
+ if isinstance(where, Filter):
77
+ if where_doc:
78
+ raise ValueError(
79
+ "Cannot use `where_document` with Filter object passed as `where` argument. "
80
+ "Please use a dictionary instead."
81
+ )
82
+ return where
83
+
71
84
  conditions = []
85
+ _and = True
86
+ if where:
87
+ if "$or" in where:
88
+ _and = False
89
+ for i in where["$or"]:
90
+ for k, v in i.items():
91
+ conditions.append(FieldCondition(key=k, match=MatchValue(value=v)))
92
+ elif "$and" in where:
93
+ _and = True
94
+ for i in where["$and"]:
95
+ for k, v in i.items():
96
+ conditions.append(FieldCondition(key=k, match=MatchValue(value=v)))
97
+ else:
98
+ for k, v in where.items():
99
+ conditions.append(FieldCondition(key=k, match=MatchValue(value=v)))
100
+
72
101
  # ChromaDB format
73
- if kwargs and "where_document" in kwargs and kwargs["where_document"]:
102
+ if where_doc:
103
+ if not _and:
104
+ raise ValueError(
105
+ "Cannot use `where_document` with `$or` condition. "
106
+ )
74
107
  conditions.append(
75
108
  FieldCondition(
76
109
  key="_text",
77
110
  match=MatchText(text=kwargs["where_document"]["$contains"])
78
111
  )
79
112
  )
80
- if where:
81
- for k, v in where.items():
82
- conditions.append(FieldCondition(key=k, match=MatchValue(value=v)))
83
113
 
84
- return Filter(must=conditions) if conditions else None
114
+ if not conditions:
115
+ return None
116
+ if _and:
117
+ return Filter(must=conditions)
118
+ return Filter(should=conditions)
85
119
 
86
120
  def search(
87
121
  self,
@@ -100,7 +100,7 @@ def _log_response(out):
100
100
  def use_logging():
101
101
  """Turns on logging of LLM requests and responses to console."""
102
102
  if not is_notebook():
103
- init(autoreset=True, strip=False)
103
+ init(strip=False)
104
104
  if _log_request not in env().llm_before_handlers:
105
105
  env().llm_before_handlers.append(_log_request)
106
106
  if _log_response not in env().llm_after_handlers:
@@ -5,6 +5,7 @@ from typing import Optional
5
5
  from dataclasses import dataclass, field
6
6
  from enum import Enum
7
7
 
8
+ import httpx
8
9
  import requests
9
10
  from fastmcp import Client
10
11
  from fastmcp.client.progress import ProgressHandler
@@ -23,6 +24,16 @@ class WrongMcpUsage(BadAIAnswer):
23
24
  ...
24
25
 
25
26
 
27
+ class HeaderAuth(httpx.Auth):
28
+ def __init__(self, header_name, header_value):
29
+ self.header_name = header_name
30
+ self.header_value = header_value
31
+
32
+ def auth_flow(self, request):
33
+ request.headers[self.header_name] = self.header_value
34
+ yield request
35
+
36
+
26
37
  class ToolsCache:
27
38
  FILE = "mcp/tools.json"
28
39
 
@@ -75,9 +86,14 @@ class MCPConnection:
75
86
  fetch_tools: bool = True,
76
87
  use_cache: bool = True,
77
88
  connect_timeout: float = 10,
89
+ auth: httpx.Auth = None,
78
90
  ) -> "MCPConnection":
79
91
  con: MCPConnection = MCPConnection(url=url, transport=transport)
80
- con._client = Client(url, timeout=connect_timeout) # pylint: disable=W0212
92
+ con._client = Client(
93
+ url,
94
+ timeout=connect_timeout,
95
+ auth=auth,
96
+ ) # pylint: disable=W0212
81
97
  await con._client.__aenter__() # pylint: disable=E1101,W0212,C2801
82
98
  if fetch_tools:
83
99
  await con.fetch_tools(use_cache=use_cache)
@@ -236,6 +252,7 @@ class MCPServer:
236
252
  name: str = field(default="")
237
253
  tools: Tools = field(default_factory=Tools)
238
254
  transport: McpTransport = field(default=None)
255
+ auth: httpx.Auth | dict | None = field(default=None)
239
256
 
240
257
  @staticmethod
241
258
  def _try_sse_or_streamable_http(url) -> tuple[McpTransport, str]:
@@ -272,6 +289,12 @@ class MCPServer:
272
289
  self.name = MCPServer.name_from_url(self.url)
273
290
  if not self.transport:
274
291
  self.transport = self._guess_transport_type_by_url(self.url)
292
+ if self.auth and isinstance(self.auth, dict) and len(self.auth)==1:
293
+ header_name, header_value = next(iter(self.auth.items()))
294
+ self.auth = HeaderAuth(
295
+ header_name=header_name,
296
+ header_value=header_value
297
+ )
275
298
 
276
299
  async def connect(
277
300
  self,
@@ -289,6 +312,7 @@ class MCPServer:
289
312
  fetch_tools=fetch_tools,
290
313
  use_cache=use_cache,
291
314
  connect_timeout=connect_timeout,
315
+ auth=self.auth
292
316
  )
293
317
 
294
318
  def get_tools_cache(self) -> Tools | None:
@@ -2,7 +2,7 @@ from colorama import Fore, init
2
2
  from .utils import is_notebook
3
3
 
4
4
  if not is_notebook():
5
- init(autoreset=True, strip=False)
5
+ init(strip=False)
6
6
 
7
7
 
8
8
  def info(*args, color=Fore.LIGHTYELLOW_EX, **kwargs):
@@ -29,14 +29,15 @@ dependencies = [
29
29
  "mcp~=1.9.2",
30
30
  "fastmcp~=2.8.0",
31
31
  "docstring_parser~=0.16.0",
32
+ "httpx~=0.28.1"
32
33
  ]
33
34
  requires-python = ">=3.10"
34
35
 
35
36
  authors = [
36
- { name = "Vitalii Stepanenko", email = "mail@vitalii.in" },
37
+ { name = "Vitalii Stepanenko", email = "mail@vitaliy.in" },
37
38
  ]
38
39
  maintainers = [
39
- { name = "Vitalii Stepanenko", email = "mail@vitalii.in" },
40
+ { name = "Vitalii Stepanenko", email = "mail@vitaliy.in" },
40
41
  ]
41
42
  license = { file = "LICENSE" }
42
43
 
File without changes