amsdal_ml 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

amsdal_ml/agents/functional_calling_agent.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator
+from typing import no_type_check
+
+from amsdal_ml.agents.agent import Agent
+from amsdal_ml.agents.agent import AgentOutput
+from amsdal_ml.agents.mcp_client_tool import ClientToolProxy
+from amsdal_ml.agents.python_tool import PythonTool
+from amsdal_ml.agents.python_tool import _PythonToolProxy
+from amsdal_ml.agents.tool_adapters import ToolAdapter
+from amsdal_ml.agents.tool_adapters import get_tool_adapter
+from amsdal_ml.fileio.base_loader import PLAIN_TEXT
+from amsdal_ml.fileio.base_loader import FileAttachment
+from amsdal_ml.mcp_client.base import ToolClient
+from amsdal_ml.mcp_client.base import ToolInfo
+from amsdal_ml.ml_models.models import MLModel
+from amsdal_ml.ml_models.models import StructuredMessage
+from amsdal_ml.ml_models.utils import ResponseFormat
+
+
+class FunctionalCallingAgent(Agent):
+    """
+    An agent that uses the native function calling capabilities of LLMs (e.g., OpenAI)
+    to execute tools and answer user queries.
+
+    Attributes:
+        model (MLModel): The LLM model instance to use (must support function calling).
+        max_steps (int): Maximum number of tool execution steps allowed.
+        per_call_timeout (float | None): Timeout in seconds for each tool execution.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: MLModel,
+        tools: list[PythonTool | ToolClient] | None = None,
+        max_steps: int = 6,
+        per_call_timeout: float | None = 20.0,
+        adapter: ToolAdapter | None = None,
+    ):
+        """
+        Initialize the FunctionalCallingAgent.
+
+        Args:
+            model: The LLM model to use.
+            tools: A list of tools (PythonTool or ToolClient) available to the agent.
+            max_steps: The maximum number of iterations (model -> tool -> model) allowed.
+            per_call_timeout: Timeout for individual tool calls.
+            adapter: Optional tool adapter. If None, will auto-detect based on LLM type.
+        """
+        self._tools: list[PythonTool | ToolClient] = tools or []
+        self._indexed_tools: dict[str, ClientToolProxy | _PythonToolProxy] = {}
+        self.model = model
+        self.model.setup()
+        self.max_steps = max_steps
+        self.per_call_timeout = per_call_timeout
+        self._is_tools_index_built = False
+        self.adapter = adapter or get_tool_adapter(model)
+        self._response_format: ResponseFormat = self._select_response_format()
+
+    async def arun(
+        self,
+        user_query: str,
+        *,
+        history: list[StructuredMessage] | None = None,
+        attachments: list[FileAttachment] | None = None,
+    ) -> AgentOutput:
+        """
+        Run the agent asynchronously to answer a user query.
+
+        This method executes the main loop:
+        1. Send query and tools to the model.
+        2. If model requests tool calls, execute them and report back.
+        3. Repeat until the model provides a final answer or max_steps is reached.
+
+        Args:
+            user_query: The question or instruction from the user.
+            history: Optional chat history to continue the conversation.
+            attachments: Optional list of files/documents to include in context.
+
+        Returns:
+            AgentOutput: The final answer and metadata about used tools.
+        """
+        if not self._is_tools_index_built:
+            await self._build_clients_index()
+
+        content = self._merge_attachments(user_query, attachments)
+        messages = history.copy() if history else []
+
+        #TODO: JSON markdown tables support for nlqretriever
+        #if self._response_format == ResponseFormat.JSON_OBJECT:
+        #    messages.append({'role': 'system', 'content': 'Please respond in json format.'})
+
+        messages.append({self.model.role_field: self.model.input_role, self.model.content_field: content})  # type: ignore[misc]
+        used_tools: list[str] = []
+        tools_schema = self.adapter.get_tools_schema(self._indexed_tools)
+
+        for _ in range(self.max_steps):
+            response_str = await self.model.ainvoke(
+                input=messages,
+                tools=tools_schema if tools_schema else None,
+                response_format=self._response_format,
+            )
+
+            try:
+                response_data = json.loads(response_str)
+            except json.JSONDecodeError:
+                response_data = {self.model.role_field: self.model.output_role, self.model.content_field: response_str}
+
+            messages.append(response_data)
+
+            content_text, tool_calls = self.adapter.parse_response(response_data)
+
+            if not tool_calls:
+                return AgentOutput(answer=content_text or '', used_tools=used_tools)
+
+            for tool_call in tool_calls:
+                function_name, arguments_str, call_id = self.adapter.get_tool_call_info(tool_call)
+
+                used_tools.append(function_name)
+
+                try:
+                    args = json.loads(arguments_str)
+                    tool = self._indexed_tools[function_name]
+                    result = await tool.run(args)
+
+                    if isinstance(result, (dict, list)):
+                        content_str = json.dumps(result, ensure_ascii=False)
+                    else:
+                        content_str = str(result)
+                except Exception as e:
+                    content_str = f'Error: {e!s}'
+
+                messages.append({
+                    self.model.role_field: self.model.tool_role,  # type: ignore[misc]
+                    self.model.tool_call_id_field: call_id,
+                    self.model.tool_name_field: function_name,
+                    self.model.content_field: content_str,
+                })
+
+        return AgentOutput(
+            answer='Agent stopped due to iteration limit.',
+            used_tools=used_tools,
+        )
+
+    @no_type_check
+    async def astream(
+        self,
+        user_query: str,
+        *,
+        history: list[StructuredMessage] | None = None,
+        attachments: list[FileAttachment] | None = None,
+    ) -> AsyncIterator[str]:
+        """
+        Stream the agent's response asynchronously.
+
+        Currently, this method buffers the full execution and yields the final answer
+        at once. True streaming of intermediate steps is not yet implemented.
+
+        Args:
+            user_query: The question or instruction from the user.
+            history: Optional chat history to continue the conversation.
+            attachments: Optional list of files/documents.
+
+        Yields:
+            str: Chunks of the final answer (currently just the full answer).
+        """
+        output = await self.arun(user_query, history=history, attachments=attachments)
+        yield output.answer
+
+    def _select_response_format(self) -> ResponseFormat:
+        """
+        Select the best response format supported by the model.
+
+        Returns:
+            ResponseFormat: PLAIN_TEXT to allow raw Markdown tables.
+        """
+        return ResponseFormat.PLAIN_TEXT
+
+    async def _build_clients_index(self) -> None:
+        """
+        Build the internal index of tools.
+
+        Iterates through the provided tools, resolving ToolClients into individual
+        callable proxies and indexing PythonTools directly.
+        """
+        self._indexed_tools.clear()
+
+        for tool in self._tools:
+            if isinstance(tool, ToolClient):
+                infos: list[ToolInfo] = await tool.list_tools()
+                for ti in infos:
+                    qname = f'{ti.alias}.{ti.name}'
+                    proxy = ClientToolProxy(
+                        client=tool,
+                        alias=ti.alias,
+                        name=ti.name,
+                        schema=ti.input_schema or {},
+                        description=ti.description or '',
+                    )
+                    proxy.set_timeout(self.per_call_timeout)
+                    self._indexed_tools[qname] = proxy
+            elif isinstance(tool, PythonTool):
+                if tool.name in self._indexed_tools:
+                    msg = f'Tool name conflict: {tool.name} is already defined.'
+                    raise ValueError(msg)
+                proxy = _PythonToolProxy(tool, timeout=self.per_call_timeout)  # type: ignore[assignment]
+                self._indexed_tools[tool.name] = proxy
+            else:
+                msg = f'Unsupported tool type: {type(tool)}'
+                raise TypeError(msg)
+
+        self._is_tools_index_built = True
+
+    def _merge_attachments(self, query: str, attachments: list[FileAttachment] | None) -> str:
+        """
+        Merge plain text attachments into the user query.
+
+        Args:
+            query: The original user query.
+            attachments: Optional list of file attachments.
+
+        Returns:
+            str: The query augmented with attachment content.
+        """
+        if not attachments:
+            return query
+        extras = [str(a.content) for a in attachments if a.type == PLAIN_TEXT]
+        if not extras:
+            return query
+        return f'{query}\n\n[ATTACHMENTS]\n' + '\n\n'.join(extras)

amsdal_ml/agents/mcp_client_tool.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from typing import Any
+
+from amsdal_ml.mcp_client.base import ToolClient
+
+
+class ClientToolProxy:
+    def __init__(
+        self,
+        client: ToolClient,
+        alias: str,
+        name: str,
+        schema: dict[str, Any],
+        description: str,
+    ):
+        self.client = client
+        self.alias = alias
+        self.name = name
+        self.qualified = f'{alias}.{name}'
+        self.parameters = schema
+        self.description = description
+        self._default_timeout: float | None = 20.0
+
+    def set_timeout(self, timeout: float | None) -> None:
+        self._default_timeout = timeout
+
+    async def run(
+        self,
+        args: dict[str, Any],
+        context: Any = None,
+        *,
+        convert_result: bool = True,
+    ) -> Any:
+        _ = (context, convert_result)
+
+        if self.parameters:
+            try:
+                import jsonschema
+
+                jsonschema.validate(instance=args, schema=self.parameters)
+            except Exception as exc:
+                msg = f'Tool input validation failed for {self.qualified}: {exc}'
+                raise ValueError(msg) from exc
+
+        return await self.client.call(self.name, args, timeout=self._default_timeout)

amsdal_ml/agents/python_tool.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from collections.abc import Coroutine
+from typing import Any
+
+from pydantic import create_model
+from pydantic.fields import FieldInfo
+
+
+class PythonTool:
+    def __init__(
+        self,
+        func: Callable[..., Coroutine[Any, Any, Any] | Any],
+        name: str,
+        description: str,
+    ):
+        if not inspect.iscoroutinefunction(func) and not inspect.isfunction(func):
+            msg = 'Tool must be a function or coroutine function'
+            raise TypeError(msg)
+
+        self.func = func
+        self.name = name
+        self.description = description
+        self.is_async = inspect.iscoroutinefunction(func)
+        self.parameters = self._build_schema()
+
+    def _build_schema(self) -> dict[str, Any]:
+        sig = inspect.signature(self.func)
+        fields: dict[str, Any] = {}
+        for param in sig.parameters.values():
+            if param.kind not in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            ):
+                continue
+
+            field_info = (
+                param.default
+                if isinstance(param.default, FieldInfo)
+                else FieldInfo(
+                    default=param.default if param.default is not inspect.Parameter.empty else ...,
+                    description=None,
+                )
+            )
+            fields[param.name] = (
+                param.annotation if param.annotation is not inspect.Parameter.empty else Any,
+                field_info,
+            )
+
+        model = create_model(f'{self.name}Input', **fields)
+        schema = model.model_json_schema()
+
+        return {
+            'type': 'object',
+            'properties': schema.get('properties', {}),
+            'required': schema.get('required', []),
+        }
+
+
+class _PythonToolProxy:
+    def __init__(self, tool: PythonTool, timeout: float | None = 20.0):
+        self.tool = tool
+        self.name = tool.name
+        self.qualified = tool.name
+        self.parameters = tool.parameters
+        self.description = tool.description
+        self._default_timeout = timeout
+
+    def set_timeout(self, timeout: float | None) -> None:
+        self._default_timeout = timeout
+
+    async def run(
+        self,
+        args: dict[str, Any],
+        context: Any = None,
+        *,
+        convert_result: bool = True,
+    ) -> Any:
+        _ = (context, convert_result)
+
+        if self.tool.is_async:
+            return await self.tool.func(**args)
+        else:
+            return self.tool.func(**args)

amsdal_ml/agents/retriever_tool.py

@@ -14,12 +14,10 @@ from amsdal_ml.ml_retrievers.openai_retriever import OpenAIRetriever
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s [%(levelname)s] %(message)s',
-    handlers=[
-        logging.FileHandler("server2.log"),
-        logging.StreamHandler(sys.stdout)
-    ]
+    handlers=[logging.FileHandler('server2.log'), logging.StreamHandler(sys.stdout)],
 )
 
+
 class RetrieverArgs(BaseModel):
     query: str = Field(..., description='User search query')
     k: int = 5
@@ -29,6 +27,7 @@ class RetrieverArgs(BaseModel):
 
 class _RetrieverSingleton:
     """Singleton holder for lazy retriever initialization."""
+
     _instance: Optional[OpenAIRetriever] = None
 
     @classmethod
@@ -46,7 +45,7 @@ async def retriever_search(
     exclude_tags: Optional[list[str]] = None,
 ) -> list[dict[str, Any]]:
     logging.info(
-        f"retriever_search called with query={query}, k={k}, include_tags={include_tags}, exclude_tags={exclude_tags}"
+        f'retriever_search called with query={query}, k={k}, include_tags={include_tags}, exclude_tags={exclude_tags}'
     )
     retriever = _RetrieverSingleton.get()
     chunks = await retriever.asimilarity_search(
@@ -55,7 +54,7 @@ async def retriever_search(
         include_tags=include_tags,
         exclude_tags=exclude_tags,
     )
-    logging.info(f"retriever_search found {len(chunks)} chunks: {chunks}")
+    logging.info(f'retriever_search found {len(chunks)} chunks: {chunks}')
     out: list[dict[str, Any]] = []
     for c in chunks:
         if hasattr(c, 'model_dump'):

amsdal_ml/agents/tool_adapters.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+from typing import Any
+
+if TYPE_CHECKING:
+    from amsdal_ml.ml_models.models import MLModel
+
+
+class ToolAdapter(ABC):
+    @abstractmethod
+    def get_tools_schema(self, tools: dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Converts indexed tools to the model-specific function calling schema.
+
+        Args:
+            tools: A dictionary mapping tool names to tool proxy objects.
+                   Tool objects are expected to have `description` and `parameters` attributes.
+
+        Returns:
+            list[dict[str, Any]]: A list of tool definitions compatible with the model's API.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def parse_response(self, response_data: dict[str, Any]) -> tuple[str | None, list[dict[str, Any]] | None]:
+        """
+        Parses the model's response to extract content and tool calls.
+
+        Args:
+            response_data: The JSON response from the model.
+
+        Returns:
+            tuple[str | None, list[dict[str, Any]] | None]: A tuple containing:
+                - content_text: The text content of the response (or None).
+                - tool_calls: A list of tool calls (or None).
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_tool_call_info(self, tool_call: dict[str, Any]) -> tuple[str, str, str]:
+        """
+        Extracts function name, arguments string, and call ID from a tool call object.
+
+        Args:
+            tool_call: A single tool call object from the parsed response.
+
+        Returns:
+            tuple[str, str, str]: A tuple containing (function_name, arguments_str, call_id).
+        """
+        raise NotImplementedError
+
+
+class OpenAIToolAdapter(ToolAdapter):
+    def get_tools_schema(self, tools: dict[str, Any]) -> list[dict[str, Any]]:
+        tools_schema = []
+        for name, tool in tools.items():
+            tools_schema.append({
+                'type': 'function',
+                'function': {
+                    'name': name,
+                    'description': tool.description,
+                    'parameters': tool.parameters,
+                },
+            })
+        return tools_schema
+
+    def parse_response(self, response_data: dict[str, Any]) -> tuple[str | None, list[dict[str, Any]] | None]:
+        content_text = response_data.get('content')
+        tool_calls = response_data.get('tool_calls')
+        return content_text, tool_calls
+
+    def get_tool_call_info(self, tool_call: dict[str, Any]) -> tuple[str, str, str]:
+        function_name = tool_call['function']['name']
+        arguments_str = tool_call['function']['arguments']
+        call_id = tool_call['id']
+        return function_name, arguments_str, call_id
+
+
+
+def get_tool_adapter(model: MLModel) -> ToolAdapter:
+    """
+    Factory function to get the appropriate ToolAdapter for a given model.
+
+    Args:
+        model: The MLModel instance.
+
+    Returns:
+        ToolAdapter: An instance of a ToolAdapter subclass.
+    """
+    model_name = model.__class__.__name__.lower()
+
+    if "openai" in model_name:
+        return OpenAIToolAdapter()
+
+    return OpenAIToolAdapter()

amsdal_ml/fileio/base_loader.py

@@ -9,8 +9,8 @@ from typing import Optional
 
 from pydantic import BaseModel
 
-PLAIN_TEXT = "plain_text"
-FILE_ID = "file_id"
+PLAIN_TEXT = 'plain_text'
+FILE_ID = 'file_id'
 
 
 class FileData(BaseModel):
@@ -39,18 +39,20 @@ class FileItem:
     @staticmethod
     def from_path(path: str, *, filedata: FileData | None = None) -> FileItem:
         # Caller is responsible for lifecycle; loaders may close after upload.
-        f = open(path, "rb")
-        return FileItem(file=f, filename=path.split("/")[-1], filedata=filedata)
+        f = open(path, 'rb')
+        return FileItem(file=f, filename=path.split('/')[-1], filedata=filedata)
 
     @staticmethod
     def from_bytes(data: bytes, *, filename: str | None = None, filedata: FileData | None = None) -> FileItem:
         import io
+
         return FileItem(file=io.BytesIO(data), filename=filename, filedata=filedata)
 
     @staticmethod
     def from_str(text: str, *, filename: str | None = None, filedata: FileData | None = None) -> FileItem:
         import io
-        return FileItem(file=io.BytesIO(text.encode("utf-8")), filename=filename, filedata=filedata)
+
+        return FileItem(file=io.BytesIO(text.encode('utf-8')), filename=filename, filedata=filedata)
 
 
 class BaseFileLoader(ABC):

amsdal_ml/fileio/openai_loader.py

@@ -18,7 +18,7 @@ from amsdal_ml.fileio.base_loader import FileItem
 
 logger = logging.getLogger(__name__)
 
-AllowedPurpose = Literal["assistants", "batch", "fine-tune", "vision", "user_data", "evals"]
+AllowedPurpose = Literal['assistants', 'batch', 'fine-tune', 'vision', 'user_data', 'evals']
 
 
 class OpenAIFileLoader(BaseFileLoader):
@@ -26,35 +26,35 @@ class OpenAIFileLoader(BaseFileLoader):
     Loader which uploads files into OpenAI Files API and returns openai_file_id.
     """
 
-    def __init__(self, client: AsyncOpenAI, *, purpose: AllowedPurpose = "assistants") -> None:
+    def __init__(self, client: AsyncOpenAI, *, purpose: AllowedPurpose = 'assistants') -> None:
         self.client = client
         self.purpose: AllowedPurpose = purpose  # mypy: Literal union, matches SDK
 
     async def _upload_one(self, file: BinaryIO, *, filename: str | None, filedata: FileData | None) -> FileAttachment:
         try:
-            if hasattr(file, "seek"):
+            if hasattr(file, 'seek'):
                 file.seek(0)
         except Exception as exc:  # pragma: no cover
-            logger.debug("seek(0) failed for %r: %s", filename or file, exc)
+            logger.debug('seek(0) failed for %r: %s', filename or file, exc)
 
         buf = file if isinstance(file, io.BytesIO) else io.BytesIO(file.read())
 
-        up = await self.client.files.create(file=(filename or "upload.bin", buf), purpose=self.purpose)
+        up = await self.client.files.create(file=(filename or 'upload.bin', buf), purpose=self.purpose)
 
         meta: dict[str, Any] = {
-            "filename": filename,
-            "provider": "openai",
-            "file": {
-                "id": up.id,
-                "bytes": getattr(up, "bytes", None),
-                "purpose": getattr(up, "purpose", self.purpose),
-                "created_at": getattr(up, "created_at", None),
-                "status": getattr(up, "status", None),
-                "status_details": getattr(up, "status_details", None),
+            'filename': filename,
+            'provider': 'openai',
+            'file': {
+                'id': up.id,
+                'bytes': getattr(up, 'bytes', None),
+                'purpose': getattr(up, 'purpose', self.purpose),
+                'created_at': getattr(up, 'created_at', None),
+                'status': getattr(up, 'status', None),
+                'status_details': getattr(up, 'status_details', None),
             },
         }
         if filedata is not None:
-            meta["filedata"] = filedata.model_dump()
+            meta['filedata'] = filedata.model_dump()
 
         return FileAttachment(type=FILE_ID, content=up.id, metadata=meta)
 
@@ -63,7 +63,6 @@ class OpenAIFileLoader(BaseFileLoader):
 
     async def load_batch(self, items: Sequence[FileItem]) -> list[FileAttachment]:
         tasks = [
-            asyncio.create_task(self._upload_one(it.file, filename=it.filename, filedata=it.filedata))
-            for it in items
+            asyncio.create_task(self._upload_one(it.file, filename=it.filename, filedata=it.filedata)) for it in items
         ]
         return await asyncio.gather(*tasks)

amsdal_ml/mcp_client/base.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any
 from typing import Protocol
+from typing import runtime_checkable
 
 
 @dataclass
@@ -13,6 +14,7 @@ class ToolInfo:
     input_schema: dict[str, Any]
 
 
+@runtime_checkable
 class ToolClient(Protocol):
     alias: str
 

amsdal_ml/mcp_client/http_client.py

@@ -42,7 +42,13 @@ class HttpClient(ToolClient):
         finally:
             await stack.aclose()
 
-    async def call(self, tool_name: str, args: dict[str, Any], *, timeout: float | None = None,) -> Any:
+    async def call(
+        self,
+        tool_name: str,
+        args: dict[str, Any],
+        *,
+        timeout: float | None = None,
+    ) -> Any:
         _ = timeout  # ARG002
         stack, s = await self._session()
         try: