ai-microcore 4.0.0.dev16__tar.gz → 4.0.0.dev18__tar.gz

{ai_microcore-4.0.0.dev16 → ai_microcore-4.0.0.dev18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-microcore
-Version: 4.0.0.dev16
+Version: 4.0.0.dev18
 Summary: # Minimalistic Foundation for AI Applications
 Keywords: llm,large language models,ai,similarity search,ai search,gpt,openai,framework,adapter
 Author-email: Vitalii Stepanenko <mail@vitalii.in>

{ai_microcore-4.0.0.dev16 → ai_microcore-4.0.0.dev18}/microcore/__init__.py

@@ -23,6 +23,7 @@ from .wrappers.llm_response_wrapper import LLMResponse
 from ._llm_functions import llm, allm, llm_parallel
 from .utils import parse, dedent
 from .metrics import Metrics
+from .interactive_setup import interactive_setup
 
 
 def tpl(file: os.PathLike[str] | str, **kwargs) -> str | PromptWrapper:
@@ -181,7 +182,8 @@ __all__ = [
     "mcp_server",
     "tokenizing",
     "Metrics",
+    "interactive_setup",
     # "wrappers",
 ]
 
-__version__ = "4.0.0-dev16"
+__version__ = "4.0.0-dev18"

{ai_microcore-4.0.0.dev16 → ai_microcore-4.0.0.dev18}/microcore/mcp.py

@@ -2,8 +2,11 @@ import asyncio
 import logging
 from typing import Optional
 from dataclasses import dataclass, field
+from enum import Enum
 
+import requests
 from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.sse import sse_client
 from mcp import ClientSession, types
 
 from .utils import ExtendedString, ConvertableToMessage
@@ -37,17 +40,32 @@ class ToolsCache:
         cached_tools[mcp_url] = tools
         storage.write_json(ToolsCache.FILE, cached_tools)
 
+    @staticmethod
+    def clear():
+        logging.info("Clearing MCP tools cache...")
+        storage.delete(ToolsCache.FILE)
+
 
 class MCPAnswer(ExtendedString, ConvertableToMessage):
     ...
 
 
+class McpTransport(str, Enum):
+    SSE: str = "sse"
+    STREAMABLE_HTTP: str = "streamable_http"
+    WS: str = "ws"
+    STDIO: str = "stdio"
+
+    def __str__(self):
+        return self.value
+
+
 @dataclass
 class MCPConnection:
     url: str = None
+    transport: McpTransport = field(default=McpTransport.STREAMABLE_HTTP)
     read_stream: any = None
     write_stream: any = None
-    connection: any = None
     context_manager: any = None
     session: ClientSession = field(default=None, init=False)
     tools: Optional["Tools"] = field(default=None, init=False)
@@ -57,8 +75,10 @@ class MCPConnection:
     @staticmethod
     async def init(
         url: str,
+        transport: McpTransport,
         fetch_tools: bool = True,
         use_cache: bool = True,
+        connect_timeout: float = 10,
     ) -> "MCPConnection":
 
         del_event = asyncio.Event()
@@ -69,17 +89,26 @@ class MCPConnection:
         # That's a bit of a hack for closing the async context managers
         async def lifecycle():
             try:
-                logging.info(f"Connecting to MCP {url}...")
-                context_manager = streamablehttp_client(url)
-                (
-                    read_stream,
-                    write_stream,
-                    connection
-                ) = await context_manager.__aenter__()  # pylint: disable=E1101
+                logging.info(f"Connecting to {transport} MCP {url}...")
+                if transport == McpTransport.STREAMABLE_HTTP:
+                    context_manager = streamablehttp_client(url)
+                    (
+                        read_stream,
+                        write_stream,
+                        _
+                    ) = await context_manager.__aenter__()  # pylint: disable=E1101
+                elif transport == McpTransport.SSE:
+                    context_manager = sse_client(url)
+                    (
+                        read_stream,
+                        write_stream
+                    ) = await context_manager.__aenter__()  # pylint: disable=E1101
+                else:
+                    raise ValueError(f"Unsupported transport type: {transport}")
                 con.url = url
+                con.transport = transport
                 con.read_stream = read_stream
                 con.write_stream = write_stream
-                con.connection = connection
                 con.context_manager = context_manager
                 await con.init_session()
                 if fetch_tools:
@@ -91,7 +120,15 @@ class MCPConnection:
                 await con._close()  # pylint: disable=W0212
 
         con._lifecycle_task = asyncio.create_task(lifecycle())  # pylint: disable=W0212
-        await opened_event.wait()
+        try:
+            await asyncio.wait_for(opened_event.wait(), timeout=connect_timeout)
+        except Exception as e:
+            logging.warning(f"Failed to connect to MCP {url}: {e}")
+            try:
+                await con._close()  # pylint: disable=W0212
+            except:  # noqa: E722 # pylint: disable=W0702
+                pass
+            raise
         return con
 
     async def close(self):
@@ -138,12 +175,24 @@ class MCPConnection:
         mcp_tools = await self.session.list_tools()
         self.tools = Tools.from_list([Tool.from_mcp(tool) for tool in mcp_tools.tools])
         if use_cache:
-            ToolsCache.write(self.url, self.tools)
+            self.update_tools_cache()
         return self.tools
 
+    def update_tools_cache(self):
+        if self.tools is None:
+            raise RuntimeError("Tools are not fetched yet. Call fetch_tools() first.")
+        ToolsCache.write(self.url, self.tools)
+
     def __del__(self):
         self._del_event.set()
 
+    async def call(self, name: str, **kwargs):
+        assert env().config.AI_SYNTAX_FUNCTION_NAME_FIELD not in kwargs
+        params = dict(kwargs)
+        params[env().config.AI_SYNTAX_FUNCTION_NAME_FIELD] = name
+        return await self.exec(params)
+
+
     async def exec(self, params: dict | LLMResponse):
         if isinstance(params, LLMResponse):
             try:
@@ -247,6 +296,32 @@ class MCPServer:
     url: str
     name: str = field(default="")
     tools: Tools = field(default_factory=Tools)
+    transport: McpTransport = field(default=None)
+
+    @staticmethod
+    def _try_sse_or_streamable_http(url) -> tuple[McpTransport, str]:
+        if url.endswith("/"):
+            test_sse_url = f"{url}sse"
+        else:
+            test_sse_url = f"{url}/sse"
+        try:
+            response = requests.request(method="HEAD", url=test_sse_url, timeout=5)
+            if response.status_code == 200:
+                return McpTransport.SSE, test_sse_url
+        except requests.RequestException:
+            # not a SSE endpoint or not reachable
+            pass
+        return McpTransport.STREAMABLE_HTTP, f"{url}/mcp"
+
+    @staticmethod
+    def _guess_transport_type_by_url(url: str) -> Optional[McpTransport]:
+        if url.startswith("ws://") or url.startswith("wss://"):
+            return McpTransport.WS
+        if url.endswith("/mcp"):
+            return McpTransport.STREAMABLE_HTTP
+        if url.endswith("/sse"):
+            return McpTransport.SSE
+        return None
 
     @staticmethod
     def name_from_url(url: str) -> str:
@@ -256,13 +331,29 @@ class MCPServer:
     def __post_init__(self):
         if not self.name:
             self.name = MCPServer.name_from_url(self.url)
+        if not self.transport:
+            self.transport = self._guess_transport_type_by_url(self.url)
 
     async def connect(
         self,
         fetch_tools: bool = True,
         use_cache: bool = True,
+        connect_timeout: float = 10,
     ) -> MCPConnection:
-        return await MCPConnection.init(self.url, fetch_tools=fetch_tools, use_cache=use_cache)
+        if self.transport:
+            transport, url = self.transport, self.url
+        else:
+            transport, url = self._try_sse_or_streamable_http(self.url)
+        return await MCPConnection.init(
+            url=url,
+            transport=transport,
+            fetch_tools=fetch_tools,
+            use_cache=use_cache,
+            connect_timeout=connect_timeout,
+        )
+
+    def get_tools_cache(self) -> Tools | None:
+        return ToolsCache.read(self.url)
 
 
 class MCPRegistry(dict[str, MCPServer]):
@@ -281,6 +372,30 @@ class MCPRegistry(dict[str, MCPServer]):
             raise ValueError(f"MCP server '{server_name}' not found in registry")
         return self[server_name]
 
+    async def precache_tools(
+        self,
+        raise_errors: bool = False,
+        connect_timeout: int = 10,
+    ):
+        async def precache_server_tools(server_name):
+            conn = None
+            try:
+                conn = await self.get(server_name).connect(
+                    fetch_tools=True,
+                    use_cache=False,
+                    connect_timeout=connect_timeout,
+                )
+                conn.update_tools_cache()
+            except Exception as e:  # pylint: disable=W0718
+                logging.error("Failed to precache tools for MCP server %s: %s", server_name, e)
+                if raise_errors:
+                    raise
+            finally:
+                if conn is not None:
+                    await conn.close()
+
+        await asyncio.gather(*[precache_server_tools(srv) for srv in self.keys()])
+
     async def connect_to(
         self,
         server_name: str,