amsdal_ml 0.1.0__py3-none-any.whl

amsdal_ml/__about__.py

@@ -0,0 +1 @@
+__version__ = '0.1.0'

amsdal_ml/agents/agent.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import AsyncIterator
+from collections.abc import Iterator
+from typing import Any
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class AgentMessage(BaseModel):
+    role: Literal["SYSTEM", "USER", "ASSISTANT"]
+    content: str
+
+
+class AgentOutput(BaseModel):
+    answer: str
+    used_tools: list[str] = Field(default_factory=list)
+    citations: list[dict[str, Any]] = Field(default_factory=list)
+
+
+class Agent(ABC):
+    @abstractmethod
+    async def arun(self, user_query: str) -> AgentOutput: ...
+    @abstractmethod
+    async def astream(self, user_query: str) -> AsyncIterator[str]:
+        """Yield streamed chunks for the given query."""
+        raise NotImplementedError
+
+    def run(self, user_query: str) -> AgentOutput:
+        msg = "This agent is async-only. Use arun()."
+        raise NotImplementedError(msg)
+
+    def stream(self, user_query: str) -> Iterator[str]:
+        msg = "This agent is async-only. Use astream()."
+        raise NotImplementedError(msg)

amsdal_ml/agents/default_qa_agent.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import inspect
+import json
+import re
+from collections.abc import AsyncIterator
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+from typing import no_type_check
+
+from amsdal_ml.agents.agent import Agent
+from amsdal_ml.agents.agent import AgentOutput
+from amsdal_ml.agents.promts import get_prompt
+from amsdal_ml.mcp_client.base import ToolClient
+from amsdal_ml.mcp_client.base import ToolInfo
+from amsdal_ml.ml_models.models import MLModel
+
+# ---------- STRICT ReAct REGEX ----------
+_TOOL_CALL_RE = re.compile(
+    r"Thought:\s*Do I need to use a tool\?\s*Yes\s*"
+    r"Action:\s*(?P<action>[^\n]+)\s*"
+    r"Action Input:\s*(?P<input>\{.*\})\s*$",
+    re.DOTALL | re.IGNORECASE,
+)
+
+_FINAL_RE = re.compile(
+    r"Thought:\s*Do I need to use a tool\?\s*No\s*"
+    r"Final Answer:\s*(?P<answer>.+)",
+    re.DOTALL | re.IGNORECASE,
+)
+
+
+
+# ---------- STRICT ReAct REGEX ----------
+
+@dataclass
+class Route:
+    name: str
+    match: Callable[[str], bool]
+    handler: Callable[[str], AgentOutput]
+
+
+class ParseErrorMode(Enum):
+    RAISE = "raise"
+    RETRY = "retry"
+
+
+# === proxy-tool ===
+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=None,
+        *,
+        convert_result: bool = True,
+    ):
+        _ = (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)
+
+class DefaultQAAgent(Agent):
+
+    def __init__(
+        self,
+        *,
+        model: MLModel,
+        tool_clients: list[ToolClient] | None = None,
+        max_steps: int = 6,
+        on_parse_error: ParseErrorMode = ParseErrorMode.RAISE,
+        enable_stop_guard: bool = True,
+        per_call_timeout: float | None = 20.0,
+    ):
+
+        # Only clients MCP (stdio/sse)
+        self._tool_clients: list[ToolClient] = tool_clients or []
+        self._indexed_tools: dict[str, Any] = {}  # qualified -> proxy
+
+        self.model = model
+        self.model.setup()
+        self.max_steps = max_steps
+        self.per_call_timeout = per_call_timeout
+        self.on_parse_error = on_parse_error
+        self.enable_stop_guard = enable_stop_guard
+
+        self._tools_index_built = False
+
+    # ---------- tools helpers ----------
+    def _get_tool(self, name: str) -> Any:
+        """
+        Look up tools ONLY among client-indexed tools.
+        Expected names are qualified: '<alias>.<tool_name>'.
+        """
+        if not self._tools_index_built:
+            msg = "Tool index not built. Ensure arun()/astream() was used."
+            raise RuntimeError(msg)
+        if name in self._indexed_tools:
+            return self._indexed_tools[name]
+        available = sorted(self._indexed_tool_names())
+        msg = f"Unknown tool: {name}. Available: {', '.join(available)}"
+        raise KeyError(msg)
+
+    def _indexed_tool_names(self) -> list[str]:
+        return list(self._indexed_tools.keys()) if self._tools_index_built else []
+
+    def _tool_names(self) -> str:
+        return ", ".join(sorted(self._indexed_tool_names()))
+
+    def _tool_descriptions(self) -> str:
+        parts: list[str] = []
+        if self._tools_index_built:
+            for qn, t in self._indexed_tools.items():
+                desc = t.description or "No description."
+                try:
+                    schema_json = json.dumps(t.parameters or {}, ensure_ascii=False)
+                except Exception:
+                    schema_json = str(t.parameters)
+                parts.append(f"- {qn}: {desc}\n  Args JSON schema: {schema_json}")
+        return "\n".join(parts)
+
+    async def _build_clients_index(self):
+        self._indexed_tools.clear()
+        for client in self._tool_clients:
+            infos: list[ToolInfo] = await client.list_tools()
+            for ti in infos:
+                qname = f"{ti.alias}.{ti.name}"
+                proxy = _ClientToolProxy(
+                    client=client,
+                    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
+        self._tools_index_built = True
+
+    # ---------- prompt composition ----------
+    def _react_text(self, user_query: str, scratchpad: str) -> str:
+        tmpl = get_prompt("react_chat")
+        return tmpl.render_text(
+            user_query=user_query,
+            tools=self._tool_descriptions(),
+            tool_names=self._tool_names(),
+            agent_scratchpad=scratchpad,
+            chat_history="",
+        )
+
+    @staticmethod
+    def _stopped_message() -> str:
+        return "Agent stopped due to iteration limit or time limit."
+
+    def _stopped_response(self, used_tools: list[str]) -> AgentOutput:
+        return AgentOutput(answer=self._stopped_message(), used_tools=used_tools, citations=[])
+
+    @staticmethod
+    def _serialize_observation(content: Any) -> str:
+        if isinstance(content, str | bytes):
+            return content if isinstance(content, str) else content.decode("utf-8", errors="ignore")
+        try:
+            return json.dumps(content, ensure_ascii=False)
+        except Exception:
+            return str(content)
+
+    # ---------- core run ----------
+    def run(self, _user_query: str) -> AgentOutput:
+        msg = "DefaultQAAgent is async-only for now. Use arun()."
+        raise NotImplementedError(msg)
+
+    async def _run_async(self, user_query: str) -> AgentOutput:
+        if self._tool_clients and not self._tools_index_built:
+            await self._build_clients_index()
+
+        scratch = ""
+        used_tools: list[str] = []
+        parse_retries = 0
+
+        for _ in range(self.max_steps):
+            prompt = self._react_text(user_query, scratch)
+            out = await self.model.ainvoke(prompt)
+
+            m_final = _FINAL_RE.search(out or "")
+            if m_final:
+                return AgentOutput(
+                    answer=(m_final.group("answer") or "").strip(),
+                    used_tools=used_tools,
+                    citations=[],
+                )
+
+            m_tool = _TOOL_CALL_RE.search(out or "")
+            if not m_tool:
+                parse_retries += 1
+                if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
+                    msg = (
+                        "Invalid ReAct output. Expected EXACT format (Final or Tool-call). "
+                        f"Got:\n{out}"
+                    )
+                    raise ValueError(
+                        msg
+                    )
+                scratch += (
+                    "\nThought: Previous output violated the strict format. "
+                    "Reply again using EXACTLY one of the two specified formats.\n"
+                )
+                continue
+
+            action = m_tool.group("action").strip()
+            raw_input = m_tool.group("input").strip()
+
+            try:
+                args = json.loads(raw_input)
+                if not isinstance(args, dict):
+                    msg = "Action Input must be a JSON object."
+                    raise ValueError(msg)
+            except Exception as e:
+                parse_retries += 1
+                if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
+                    msg = f"Invalid Action Input JSON: {raw_input!r} ({e})"
+                    raise ValueError(msg) from e
+                scratch += (
+                    "\nThought: Action Input must be a ONE-LINE JSON object. "
+                    "Retry with correct JSON.\n"
+                )
+                continue
+
+            tool = self._get_tool(action)
+
+            try:
+                result = await tool.run(args, context=None, convert_result=True)
+            except Exception as e:
+                # unified error payload
+                err = {
+                    "error": {
+                        "type": e.__class__.__name__,
+                        "server": getattr(tool, "alias", "local"),
+                        "tool": getattr(tool, "name", getattr(tool, "qualified", "unknown")),
+                        "message": str(e),
+                        "retryable": False,
+                    }
+                }
+                result = err
+
+            used_tools.append(action)
+            observation = self._serialize_observation(result)
+
+            scratch += (
+                "\nThought: Do I need to use a tool? Yes"
+                f"\nAction: {action}"
+                f"\nAction Input: {raw_input}"
+                f"\nObservation: {observation}\n"
+            )
+
+        return self._stopped_response(used_tools)
+
+    # ---------- public APIs ----------
+    async def arun(self, user_query: str) -> AgentOutput:
+        return await self._run_async(user_query)
+
+    # ---------- streaming ----------
+    @no_type_check
+    async def astream(self, user_query: str) -> AsyncIterator[str]:
+        if self._tool_clients and not self._tools_index_built:
+            await self._build_clients_index()
+
+        scratch = ""
+        used_tools: list[str] = []
+        parse_retries = 0
+
+        for _ in range(self.max_steps):
+            prompt = self._react_text(user_query, scratch)
+
+            buffer = ""
+
+            # Normalize model.astream: it might be an async iterator already,
+            # or a coroutine (or nested coroutines) that resolves to one.
+            _val = self.model.astream(prompt)
+            while inspect.iscoroutine(_val):
+                _val = await _val
+
+            # Optional guard (helpful during tests)
+            if not hasattr(_val, "__aiter__"):
+                msg = f"model.astream() did not yield an AsyncIterator; got {type(_val)!r}"
+                raise TypeError(msg)
+
+            model_stream = _val  # now an AsyncIterator[str]
+
+            async for chunk in model_stream:
+                buffer += chunk
+
+            m_final = _FINAL_RE.search(buffer or "")
+            if m_final:
+                answer = (m_final.group("answer") or "").strip()
+                if answer:
+                    yield answer
+                return
+
+            m_tool = _TOOL_CALL_RE.search(buffer or "")
+            if not m_tool:
+                parse_retries += 1
+                if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
+                    msg = f"Invalid ReAct output (stream). Expected EXACT format. Got:\n{buffer}"
+                    raise ValueError(msg)
+                scratch += (
+                    "\nThought: Previous output violated the strict format. "
+                    "Reply again using EXACTLY one of the two specified formats.\n"
+                )
+                continue
+
+            action = m_tool.group("action").strip()
+            raw_input = m_tool.group("input").strip()
+
+            try:
+                args = json.loads(raw_input)
+                if not isinstance(args, dict):
+                    msg = "Action Input must be a JSON object."
+                    raise ValueError(msg)
+            except Exception as e:
+                parse_retries += 1
+                if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
+                    msg = f"Invalid Action Input JSON: {raw_input!r} ({e})"
+                    raise ValueError(msg) from e
+                scratch += (
+                    "\nThought: Action Input must be a ONE-LINE JSON object. "
+                    "Retry with correct JSON.\n"
+                )
+                continue
+
+            tool = self._get_tool(action)
+
+            try:
+                result = await tool.run(args, context=None, convert_result=True)
+            except Exception as e:
+                result = {
+                    "error": {
+                        "type": e.__class__.__name__,
+                        "server": getattr(tool, "alias", "local"),
+                        "tool": getattr(tool, "name", getattr(tool, "qualified", "unknown")),
+                        "message": str(e),
+                        "retryable": False,
+                    }
+                }
+
+            used_tools.append(action)
+            observation = self._serialize_observation(result)
+
+            scratch += (
+                "\nThought: Do I need to use a tool? Yes"
+                f"\nAction: {action}"
+                f"\nAction Input: {raw_input}"
+                f"\nObservation: {observation}\n"
+            )
+
+        yield self._stopped_message()

amsdal_ml/agents/promts/__init__.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from typing import Final
+
+# avoid "magic numbers"
+_PARTS_EXPECTED: Final[int] = 2
+
+
+class _SafeDict(dict[str, Any]):
+    def __missing__(self, key: str) -> str:
+        return "{" + key + "}"
+
+
+@dataclass
+class Prompt:
+    name: str
+    system: str
+    user: str
+
+    def render_text(self, **kwargs: Any) -> str:
+        data = _SafeDict(**kwargs)
+        sys_txt = self.system.format_map(data)
+        usr_txt = self.user.format_map(data)
+        return f"{sys_txt}\n\n{usr_txt}".strip()
+
+    def render_messages(self, **kwargs: Any) -> list[dict[str, str]]:
+        data = _SafeDict(**kwargs)
+        return [
+            {"role": "system", "content": self.system.format_map(data)},
+            {"role": "user", "content": self.user.format_map(data)},
+        ]
+
+
+_prompt_cache: dict[str, Prompt] = {}
+
+
+def _load_file(name: str) -> Prompt:
+    base = Path(__file__).resolve().parent
+    path = base / f"{name}.prompt"
+    if not path.exists():
+        msg = f"Prompt '{name}' not found at {path}"
+        raise FileNotFoundError(msg)
+    raw = path.read_text(encoding="utf-8")
+    parts = raw.split("\n---\n", 1)
+    if len(parts) == _PARTS_EXPECTED:
+        system, user = parts
+    else:
+        system, user = raw, "{input}"
+    return Prompt(name=name, system=system.strip(), user=user.strip())
+
+
+def get_prompt(name: str) -> Prompt:
+    if name not in _prompt_cache:
+        _prompt_cache[name] = _load_file(name)
+    return _prompt_cache[name]

amsdal_ml/agents/promts/react_chat.prompt

@@ -0,0 +1,37 @@
+You are a concise, helpful assistant. Use tools only when they improve accuracy.
+
+TOOLS
+------
+{tools}
+
+FORMAT (STRICT — EXACTLY ONE; NO extra lines, NO Markdown)
+----------------------------------------------------------
+1. If you do NOT need a tool:
+Thought: Do I need to use a tool? No
+Final Answer: <your final answer in plain text>
+
+2. If you DO need a tool:
+Thought: Do I need to use a tool? Yes
+Action: <one tool name from: {tool_names}>
+Action Input: <ONE-LINE JSON object with the tool arguments>
+
+RULES
+-----
+- Output MUST match exactly one of the two blocks above.
+- `Action Input` MUST be a valid ONE-LINE JSON object (e.g. {{"a": 1, "b": 2}}).
+- Do NOT add anything before/after the block.
+- Do NOT print "Observation". The system will add it after tool execution.
+
+PREVIOUS CONVERSATION
+--------------------
+{chat_history}
+
+NEW INPUT
+---------
+{user_query}
+
+SCRATCHPAD
+----------
+{agent_scratchpad}
+
+Assistant:

amsdal_ml/agents/retriever_tool.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Any
+from typing import Optional
+
+from mcp.server.fastmcp.tools.base import Tool
+from pydantic import BaseModel
+from pydantic import Field
+
+from amsdal_ml.ml_retrievers.openai_retriever import OpenAIRetriever
+
+
+class RetrieverArgs(BaseModel):
+    query: str = Field(..., description='User search query')
+    k: int = 5
+    include_tags: Optional[list[str]] = None
+    exclude_tags: Optional[list[str]] = None
+
+
+_retriever = OpenAIRetriever()
+
+
+async def retriever_search(args: RetrieverArgs) -> list[dict[str, Any]]:
+    chunks = await _retriever.asimilarity_search(
+        query=args.query,
+        k=args.k,
+        include_tags=args.include_tags,
+        exclude_tags=args.exclude_tags,
+    )
+    out: list[dict[str, Any]] = []
+    for c in chunks:
+        if hasattr(c, 'model_dump'):
+            out.append(c.model_dump())
+        elif hasattr(c, 'dict'):
+            out.append(c.dict())
+        elif isinstance(c, dict):
+            out.append(c)
+        else:
+            out.append({'raw_text': str(c)})
+    return out
+
+
+retriever_tool = Tool.from_function(
+    retriever_search,
+    name='retriever.search',
+    description='Semantic search in knowledge base (OpenAI embeddings)',
+    structured_output=True,
+)

amsdal_ml/app.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+# During type checking, point to the known-available fallback.
+if TYPE_CHECKING:
+    from amsdal.contrib.app_config import AppConfig as BaseAppConfig
+else:
+    try:
+        # At runtime prefer the real AMSDAL core if present.
+        from amsdal.configs.app import AppConfig as BaseAppConfig  # type: ignore[import-not-found]
+    except Exception:  # pragma: no cover
+        from amsdal.contrib.app_config import AppConfig as BaseAppConfig
+
+
+class MLPluginAppConfig(BaseAppConfig):
+    name = "amsdal_ml"
+    verbose_name = "AMSDAL ML Plugin"
+
+    def on_ready(self) -> None:
+        pass
+
+    def on_server_startup(self) -> None:
+        pass

amsdal_ml/mcp_client/base.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+from typing import Protocol
+
+
+@dataclass
+class ToolInfo:
+    alias: str
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+class ToolClient(Protocol):
+    alias: str
+
+    async def list_tools(self) -> list[ToolInfo]: ...
+    async def call(self, tool_name: str, args: dict[str, Any], *, timeout: float | None = None) -> Any: ...

amsdal_ml/mcp_client/http_client.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from contextlib import AsyncExitStack
+from typing import Any
+from typing import Optional
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+
+from amsdal_ml.mcp_client.base import ToolClient
+from amsdal_ml.mcp_client.base import ToolInfo
+
+
+class HttpClient(ToolClient):
+    def __init__(self, *, alias: str, url: str, headers: Optional[dict[str, str]] = None):
+        self.alias = alias
+        self.url = url
+        self.headers = headers or {}
+
+    async def _session(self):
+        stack = AsyncExitStack()
+        rx, tx = await stack.enter_async_context(sse_client(self.url, headers=self.headers))
+        s = await stack.enter_async_context(ClientSession(rx, tx))
+        await s.initialize()
+        return stack, s
+
+    async def list_tools(self) -> list[ToolInfo]:
+        stack, s = await self._session()
+        try:
+            resp = await s.list_tools()
+            out: list[ToolInfo] = []
+            for t in resp.tools:
+                out.append(
+                    ToolInfo(
+                        alias=self.alias,
+                        name=t.name,
+                        description=t.description or '',
+                        input_schema=(getattr(t, 'inputSchema', None) or {}),
+                    )
+                )
+            return out
+        finally:
+            await stack.aclose()
+
+    async def call(self, tool_name: str, args: dict[str, Any], *, timeout: float | None = None,) -> Any:
+        _ = timeout  # ARG002
+        stack, s = await self._session()
+        try:
+            res = await s.call_tool(tool_name, args)
+            return getattr(res, 'content', res)
+        finally:
+            await stack.aclose()