loom-agent 0.0.1__py3-none-any.whl

loom/llm/pool.py

@@ -0,0 +1,170 @@
+"""模型池管理系统
+
+框架级的多模型管理机制，允许用户：
+1. 配置多个 LLM 模型
+2. 为每个模型设置别名和能力
+3. Agent 根据任务需求自动选择合适的模型
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Callable, Any
+from dataclasses import dataclass
+from loom.interfaces.llm import BaseLLM
+from .config import LLMCapabilities
+
+
+@dataclass
+class ModelEntry:
+    name: str
+    llm: BaseLLM
+    capabilities: LLMCapabilities
+    metadata: Dict[str, Any] = None
+
+    def __post_init__(self):
+        if self.metadata is None:
+            self.metadata = {}
+
+
+class ModelPool:
+    def __init__(self, default_model: Optional[str] = None):
+        self._models: Dict[str, ModelEntry] = {}
+        self._default_model = default_model
+
+    def add(
+        self,
+        name: str,
+        llm: BaseLLM,
+        capabilities: Optional[LLMCapabilities] = None,
+        **metadata
+    ) -> None:
+        if capabilities is None:
+            if hasattr(llm, 'capabilities'):
+                capabilities = llm.capabilities  # type: ignore[attr-defined]
+            else:
+                capabilities = LLMCapabilities()
+
+        entry = ModelEntry(
+            name=name,
+            llm=llm,
+            capabilities=capabilities,
+            metadata=metadata
+        )
+
+        self._models[name] = entry
+        if len(self._models) == 1 and self._default_model is None:
+            self._default_model = name
+
+    def get(self, name: str) -> Optional[BaseLLM]:
+        entry = self._models.get(name)
+        return entry.llm if entry else None
+
+    def get_default(self) -> Optional[BaseLLM]:
+        if self._default_model:
+            return self.get(self._default_model)
+        return None
+
+    def set_default(self, name: str) -> None:
+        if name not in self._models:
+            raise ValueError(f"Model '{name}' not found in pool")
+        self._default_model = name
+
+    def select_by_capabilities(
+        self,
+        required_capabilities: LLMCapabilities,
+        prefer_default: bool = False
+    ) -> Optional[BaseLLM]:
+        if prefer_default and self._default_model:
+            default_entry = self._models[self._default_model]
+            if self._capabilities_match(default_entry.capabilities, required_capabilities):
+                return default_entry.llm
+
+        for entry in self._models.values():
+            if self._capabilities_match(entry.capabilities, required_capabilities):
+                return entry.llm
+        return None
+
+    def select(
+        self,
+        selector: Callable[[Dict[str, ModelEntry]], Optional[str]] = None,
+        **requirements
+    ) -> Optional[BaseLLM]:
+        if selector:
+            selected_name = selector(self._models)
+            return self.get(selected_name) if selected_name else None
+
+        if requirements:
+            required_caps = LLMCapabilities(**requirements)
+            return self.select_by_capabilities(required_caps)
+        return self.get_default()
+
+    def list_models(self) -> Dict[str, ModelEntry]:
+        return dict(self._models)
+
+    def remove(self, name: str) -> None:
+        if name in self._models:
+            del self._models[name]
+            if self._default_model == name:
+                self._default_model = next(iter(self._models), None)
+
+    def __len__(self) -> int:
+        return len(self._models)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._models
+
+    @staticmethod
+    def _capabilities_match(
+        model_caps: LLMCapabilities,
+        required_caps: LLMCapabilities
+    ) -> bool:
+        if required_caps.supports_tools and not model_caps.supports_tools:
+            return False
+        if required_caps.supports_vision and not model_caps.supports_vision:
+            return False
+        if required_caps.supports_json_mode and not model_caps.supports_json_mode:
+            return False
+        if required_caps.max_tokens > model_caps.max_tokens:
+            return False
+        if required_caps.context_window > model_caps.context_window:
+            return False
+        return True
+
+
+class CapabilityBasedSelector:
+    def __init__(
+        self,
+        prefer_default: bool = True,
+        fallback_to_default: bool = True
+    ):
+        self.prefer_default = prefer_default
+        self.fallback_to_default = fallback_to_default
+
+    def select(
+        self,
+        pool: ModelPool,
+        task_context: Optional[Dict[str, Any]] = None
+    ) -> Optional[BaseLLM]:
+        if not task_context:
+            return pool.get_default()
+
+        required_caps = self._infer_capabilities(task_context)
+        selected = pool.select_by_capabilities(
+            required_caps,
+            prefer_default=self.prefer_default
+        )
+        if selected is None and self.fallback_to_default:
+            selected = pool.get_default()
+        return selected
+
+    def _infer_capabilities(
+        self,
+        task_context: Dict[str, Any]
+    ) -> LLMCapabilities:
+        return LLMCapabilities(
+            supports_vision=task_context.get("has_image", False),
+            supports_tools=task_context.get("needs_tools", False),
+            context_window=task_context.get("context_size", 8192),
+            max_tokens=task_context.get("max_output", 4096)
+        )
+

loom/llm/registry.py

@@ -0,0 +1,201 @@
+"""模型能力注册表
+
+维护各个 LLM 提供商和模型的能力信息，用于能力检测与查询。
+"""
+
+from __future__ import annotations
+
+from typing import Dict
+from .config import LLMCapabilities, LLMProvider
+
+
+class ModelRegistry:
+    OPENAI_MODELS: Dict[str, LLMCapabilities] = {
+        "gpt-4": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=False,
+            supports_streaming=True,
+            supports_json_mode=True,
+            supports_system_message=True,
+            max_tokens=8192,
+            context_window=8192,
+        ),
+        "gpt-4-turbo": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=True,
+            supports_streaming=True,
+            supports_json_mode=True,
+            supports_system_message=True,
+            max_tokens=4096,
+            context_window=128000,
+        ),
+        "gpt-4o": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=True,
+            supports_streaming=True,
+            supports_json_mode=True,
+            supports_system_message=True,
+            max_tokens=16384,
+            context_window=128000,
+        ),
+        "gpt-4o-mini": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=True,
+            supports_streaming=True,
+            supports_json_mode=True,
+            supports_system_message=True,
+            max_tokens=16384,
+            context_window=128000,
+        ),
+        "gpt-3.5-turbo": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=False,
+            supports_streaming=True,
+            supports_json_mode=True,
+            supports_system_message=True,
+            max_tokens=4096,
+            context_window=16385,
+        ),
+    }
+
+    ANTHROPIC_MODELS: Dict[str, LLMCapabilities] = {
+        "claude-3-5-sonnet-20241022": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=True,
+            supports_streaming=True,
+            supports_json_mode=False,
+            supports_system_message=True,
+            max_tokens=8192,
+            context_window=200000,
+        ),
+        "claude-3-sonnet-20240229": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=True,
+            supports_streaming=True,
+            supports_json_mode=False,
+            supports_system_message=True,
+            max_tokens=4096,
+            context_window=200000,
+        ),
+    }
+
+    GOOGLE_MODELS: Dict[str, LLMCapabilities] = {
+        "gemini-pro": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=False,
+            supports_streaming=True,
+            supports_json_mode=False,
+            supports_system_message=True,
+            max_tokens=8192,
+            context_window=32760,
+        ),
+        "gemini-1.5-pro": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=True,
+            supports_streaming=True,
+            supports_json_mode=False,
+            supports_system_message=True,
+            max_tokens=8192,
+            context_window=1000000,
+        ),
+    }
+
+    COHERE_MODELS: Dict[str, LLMCapabilities] = {
+        "command-r-plus": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=False,
+            supports_streaming=True,
+            supports_json_mode=False,
+            supports_system_message=True,
+            max_tokens=4096,
+            context_window=128000,
+        ),
+        "command-r": LLMCapabilities(
+            supports_tools=True,
+            supports_vision=False,
+            supports_streaming=True,
+            supports_json_mode=False,
+            supports_system_message=True,
+            max_tokens=4096,
+            context_window=128000,
+        ),
+    }
+
+    DEFAULT_CAPABILITIES = LLMCapabilities(
+        supports_tools=False,
+        supports_vision=False,
+        supports_streaming=True,
+        supports_json_mode=False,
+        supports_system_message=True,
+        max_tokens=4096,
+        context_window=8192,
+    )
+
+    @classmethod
+    def get_capabilities(cls, provider: str, model_name: str) -> LLMCapabilities:
+        provider = provider.lower()
+
+        if provider == LLMProvider.OPENAI or provider == "azure_openai":
+            return cls._get_with_fuzzy_match(cls.OPENAI_MODELS, model_name)
+        elif provider == LLMProvider.ANTHROPIC:
+            return cls._get_with_fuzzy_match(cls.ANTHROPIC_MODELS, model_name)
+        elif provider == LLMProvider.GOOGLE:
+            return cls._get_with_fuzzy_match(cls.GOOGLE_MODELS, model_name)
+        elif provider == LLMProvider.COHERE:
+            return cls._get_with_fuzzy_match(cls.COHERE_MODELS, model_name)
+        elif provider == LLMProvider.OLLAMA:
+            return LLMCapabilities(
+                supports_tools=False,
+                supports_vision=False,
+                supports_streaming=True,
+                supports_json_mode=False,
+                supports_system_message=True,
+                max_tokens=2048,
+                context_window=8192,
+            )
+        else:
+            return cls.DEFAULT_CAPABILITIES
+
+    @classmethod
+    def supports_tools(cls, provider: str, model_name: str) -> bool:
+        return cls.get_capabilities(provider, model_name).supports_tools
+
+    @classmethod
+    def supports_vision(cls, provider: str, model_name: str) -> bool:
+        return cls.get_capabilities(provider, model_name).supports_vision
+
+    @classmethod
+    def get_context_window(cls, provider: str, model_name: str) -> int:
+        return cls.get_capabilities(provider, model_name).context_window
+
+    @classmethod
+    def register_model(
+        cls,
+        provider: str,
+        model_name: str,
+        capabilities: LLMCapabilities
+    ):
+        provider = provider.lower()
+        if provider == LLMProvider.OPENAI:
+            cls.OPENAI_MODELS[model_name] = capabilities
+        elif provider == LLMProvider.ANTHROPIC:
+            cls.ANTHROPIC_MODELS[model_name] = capabilities
+        elif provider == LLMProvider.GOOGLE:
+            cls.GOOGLE_MODELS[model_name] = capabilities
+        elif provider == LLMProvider.COHERE:
+            cls.COHERE_MODELS[model_name] = capabilities
+
+    @classmethod
+    def _get_with_fuzzy_match(
+        cls,
+        models: Dict[str, LLMCapabilities],
+        model_name: str
+    ) -> LLMCapabilities:
+        if model_name in models:
+            return models[model_name]
+        model_name_lower = model_name.lower()
+        for key, value in models.items():
+            if model_name_lower in key.lower() or key.lower() in model_name_lower:
+                return value
+        return cls.DEFAULT_CAPABILITIES
+

loom/mcp/__init__.py

@@ -0,0 +1,4 @@
+from .registry import MCPToolRegistry
+
+__all__ = ["MCPToolRegistry"]
+

loom/mcp/client.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class MCPServerConfig:
+    command: str
+    args: List[str]
+    env: Optional[Dict[str, str]] = None
+
+
+class MCPError(Exception):
+    pass
+
+
+class MCPClient:
+    """最小 MCP 客户端：JSON-RPC over stdio。"""
+
+    def __init__(self, server_config: MCPServerConfig) -> None:
+        self.config = server_config
+        self.process: Optional[asyncio.subprocess.Process] = None
+        self.request_id = 0
+        self._response_futures: Dict[int, asyncio.Future] = {}
+
+    async def connect(self) -> None:
+        self.process = await asyncio.create_subprocess_exec(
+            self.config.command,
+            *self.config.args,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=self.config.env,
+        )
+        asyncio.create_task(self._read_responses())
+        await self._send_request(
+            "initialize",
+            {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {"name": "loom-framework", "version": "0.1.0"},
+            },
+        )
+
+    async def disconnect(self) -> None:
+        if self.process:
+            self.process.terminate()
+            await self.process.wait()
+
+    async def list_tools(self) -> List[Dict]:
+        result = await self._send_request("tools/list", {})
+        return result.get("tools", [])
+
+    async def call_tool(self, name: str, arguments: Dict) -> Dict:
+        return await self._send_request("tools/call", {"name": name, "arguments": arguments})
+
+    async def _send_request(self, method: str, params: Dict) -> Dict:
+        assert self.process and self.process.stdin and self.process.stdout
+        self.request_id += 1
+        req = {"jsonrpc": "2.0", "id": self.request_id, "method": method, "params": params}
+        fut: asyncio.Future = asyncio.get_event_loop().create_future()
+        self._response_futures[self.request_id] = fut
+        self.process.stdin.write((json.dumps(req) + "\n").encode())
+        await self.process.stdin.drain()
+        return await fut
+
+    async def _read_responses(self) -> None:
+        assert self.process and self.process.stdout
+        while True:
+            line = await self.process.stdout.readline()
+            if not line:
+                break
+            resp = json.loads(line.decode())
+            if "id" in resp:
+                rid = resp["id"]
+                fut = self._response_futures.pop(rid, None)
+                if fut is None:
+                    continue
+                if "error" in resp:
+                    fut.set_exception(MCPError(str(resp["error"])))
+                else:
+                    fut.set_result(resp.get("result", {}))
+

loom/mcp/registry.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from .client import MCPClient, MCPServerConfig
+from .tool_adapter import MCPTool
+
+
+class MCPToolRegistry:
+    """MCP 工具注册中心 - 负责发现/连接 server 并加载工具。"""
+
+    def __init__(self, config_path: Optional[str] = None) -> None:
+        self.config_path = Path(config_path) if config_path else Path.home() / ".loom" / "mcp.json"
+        self.servers: Dict[str, MCPClient] = {}
+        self.tools: Dict[str, MCPTool] = {}
+
+    async def discover_local_servers(self) -> None:
+        if not self.config_path.exists():
+            return
+        config = json.loads(self.config_path.read_text())
+        for name, scfg in config.get("mcpServers", {}).items():
+            await self.add_server(name, scfg)
+
+    async def add_server(self, name: str, cfg: Dict) -> None:
+        client = MCPClient(
+            MCPServerConfig(command=cfg["command"], args=cfg.get("args", []), env=cfg.get("env"))
+        )
+        await client.connect()
+        self.servers[name] = client
+        await self._load_server_tools(name, client)
+
+    async def _load_server_tools(self, server_name: str, client: MCPClient) -> None:
+        for spec in await client.list_tools():
+            tool = MCPTool(mcp_tool_spec=spec, mcp_client=client)
+            self.tools[f"{server_name}:{tool.name}"] = tool
+
+    async def load_from_server(self, server_name: str) -> List[MCPTool]:
+        return [tool for key, tool in self.tools.items() if key.startswith(f"{server_name}:")]
+
+    async def load_servers(self, server_names: List[str]) -> List[MCPTool]:
+        tools: List[MCPTool] = []
+        for s in server_names:
+            tools.extend(await self.load_from_server(s))
+        return tools
+
+    async def list_all_tools(self) -> Dict[str, List[str]]:
+        grouped: Dict[str, List[str]] = {}
+        for key in self.tools.keys():
+            server, tname = key.split(":", 1)
+            grouped.setdefault(server, []).append(tname)
+        return grouped
+
+    async def close(self) -> None:
+        for client in self.servers.values():
+            await client.disconnect()
+

loom/mcp/tool_adapter.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from pydantic import BaseModel, Field, create_model
+
+from loom.interfaces.tool import BaseTool
+from .client import MCPClient
+
+
+class MCPTool(BaseTool):
+    """MCP 工具适配为 Loom 工具。"""
+
+    def __init__(self, mcp_tool_spec: Dict, mcp_client: MCPClient) -> None:
+        self.mcp_spec = mcp_tool_spec
+        self.mcp_client = mcp_client
+        self.name = mcp_tool_spec["name"]
+        self.description = mcp_tool_spec.get("description", "")
+        self.args_schema = self._build_pydantic_schema(mcp_tool_spec.get("inputSchema", {}))
+
+    def _build_pydantic_schema(self, json_schema: Dict) -> type[BaseModel]:
+        properties = json_schema.get("properties", {})
+        required = set(json_schema.get("required", []))
+        fields: Dict[str, tuple[type, Any]] = {}
+        for fname, spec in properties.items():
+            py_type = self._json_type_to_python(spec.get("type", "string"))
+            desc = spec.get("description", "")
+            default = ... if fname in required else spec.get("default", None)
+            fields[fname] = (py_type, Field(default, description=desc))
+        return create_model(f"{self.name.title()}Args", **fields)  # type: ignore[arg-type]
+
+    def _json_type_to_python(self, json_type: str) -> type:
+        return {
+            "string": str,
+            "integer": int,
+            "number": float,
+            "boolean": bool,
+            "array": list,
+            "object": dict,
+        }.get(json_type, str)
+
+    async def run(self, **kwargs) -> Any:
+        result = await self.mcp_client.call_tool(self.name, kwargs)
+        content = result.get("content", [])
+        if isinstance(content, list) and content:
+            return content[0].get("text", "")
+        return result
+

loom/observability/__init__.py

@@ -0,0 +1,5 @@
+"""Observability subsystem - correlation IDs, structured logging, system reminders"""
+
+from __future__ import annotations
+
+__all__: list[str] = []

loom/patterns/__init__.py

@@ -0,0 +1,5 @@
+"""高级模式 - Multi-Agent, RAG, ReAct 等"""
+
+from loom.patterns.multi_agent import MultiAgentSystem
+
+__all__ = ["MultiAgentSystem"]