loom-agent 0.0.1__py3-none-any.whl

loom/builtin/embeddings/sentence_transformers_embedding.py

@@ -0,0 +1,145 @@
+"""Sentence Transformers Embedding 适配器（本地模型）"""
+
+from __future__ import annotations
+
+from typing import List
+
+from loom.interfaces.embedding import BaseEmbedding
+
+try:
+    from sentence_transformers import SentenceTransformer
+    SENTENCE_TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    SENTENCE_TRANSFORMERS_AVAILABLE = False
+
+
+class SentenceTransformersEmbedding(BaseEmbedding):
+    """
+    Sentence Transformers 本地 Embedding 适配器
+
+    特点:
+    - ✅ 完全本地运行，无需 API
+    - ✅ 支持多种预训练模型
+    - ✅ 支持多语言
+    - ✅ GPU 加速支持
+
+    推荐模型:
+    - all-MiniLM-L6-v2 (384 维, 快速)
+    - all-mpnet-base-v2 (768 维, 平衡)
+    - paraphrase-multilingual-MiniLM-L12-v2 (384 维, 多语言)
+
+    示例:
+        from loom.builtin.embeddings import SentenceTransformersEmbedding
+
+        # 英文模型
+        embedding = SentenceTransformersEmbedding(
+            model_name="all-MiniLM-L6-v2"
+        )
+
+        # 多语言模型
+        embedding = SentenceTransformersEmbedding(
+            model_name="paraphrase-multilingual-MiniLM-L12-v2"
+        )
+
+        # 单个文本
+        vector = await embedding.embed_query("Hello world")
+
+        # 批量文本
+        vectors = await embedding.embed_documents([
+            "Hello world",
+            "Loom framework"
+        ])
+    """
+
+    def __init__(
+        self,
+        model_name: str = "all-MiniLM-L6-v2",
+        device: str = "cpu",
+        batch_size: int = 32,
+        normalize_embeddings: bool = True,
+    ):
+        """
+        Parameters:
+            model_name: 模型名称（HuggingFace 模型 ID）
+            device: 设备 ('cpu', 'cuda', 'mps')
+            batch_size: 批处理大小
+            normalize_embeddings: 是否归一化向量
+        """
+        if not SENTENCE_TRANSFORMERS_AVAILABLE:
+            raise ImportError(
+                "Sentence Transformers is not installed. "
+                "Install with: pip install sentence-transformers"
+            )
+
+        self.model_name = model_name
+        self.device = device
+        self.batch_size = batch_size
+        self.normalize_embeddings = normalize_embeddings
+
+        # 加载模型
+        self.model = SentenceTransformer(model_name, device=device)
+
+    async def embed_query(self, text: str) -> List[float]:
+        """
+        对单个查询文本进行向量化
+
+        Parameters:
+            text: 查询文本
+
+        Returns:
+            向量（列表）
+        """
+        # Sentence Transformers 的 encode 是同步的
+        # 在 async 函数中直接调用（小型模型通常很快）
+        embedding = self.model.encode(
+            text,
+            normalize_embeddings=self.normalize_embeddings,
+            convert_to_numpy=True
+        )
+        return embedding.tolist()
+
+    async def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        批量向量化文档
+
+        Parameters:
+            texts: 文本列表
+
+        Returns:
+            向量列表
+        """
+        # 过滤空文本
+        non_empty_texts = [t for t in texts if t.strip()]
+
+        if not non_empty_texts:
+            dimension = self.get_dimension()
+            return [[0.0] * dimension] * len(texts)
+
+        # 批量编码
+        embeddings = self.model.encode(
+            non_empty_texts,
+            batch_size=self.batch_size,
+            normalize_embeddings=self.normalize_embeddings,
+            convert_to_numpy=True,
+            show_progress_bar=False
+        )
+
+        vectors = embeddings.tolist()
+
+        # 处理空文本位置
+        result = []
+        non_empty_idx = 0
+        dimension = len(vectors[0])
+
+        for text in texts:
+            if text.strip():
+                result.append(vectors[non_empty_idx])
+                non_empty_idx += 1
+            else:
+                result.append([0.0] * dimension)
+
+        return result
+
+    def get_dimension(self) -> int:
+        """返回向量维度"""
+        return self.model.get_sentence_embedding_dimension()

loom/builtin/llms/__init__.py

@@ -0,0 +1,8 @@
+from .mock import MockLLM
+from .rule import RuleLLM
+
+try:
+    from .openai import OpenAILLM
+    __all__ = ["MockLLM", "RuleLLM", "OpenAILLM"]
+except ImportError:
+    __all__ = ["MockLLM", "RuleLLM"]

loom/builtin/llms/mock.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import asyncio
+from typing import AsyncGenerator, Dict, List, Optional
+
+from loom.interfaces.llm import BaseLLM
+
+
+class MockLLM(BaseLLM):
+    """用于测试与示例的 Mock LLM。"""
+
+    def __init__(self, responses: Optional[List[str]] = None, name: str = "mock-llm") -> None:
+        self._responses = list(responses or ["OK"])
+        self._model_name = name
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    async def generate(self, messages: List[Dict]) -> str:
+        if self._responses:
+            return self._responses.pop(0)
+        return "".join(m.get("content", "") for m in messages if m.get("role") == "user")
+
+    async def stream(self, messages: List[Dict]) -> AsyncGenerator[str, None]:
+        text = await self.generate(messages)
+        # 简单逐字符流
+        for ch in text:
+            await asyncio.sleep(0)  # 让出循环
+            yield ch
+
+    async def generate_with_tools(self, messages: List[Dict], tools: List[Dict]) -> Dict:
+        return {"content": await self.generate(messages), "tool_calls": []}
+

loom/builtin/llms/openai.py

@@ -0,0 +1,168 @@
+"""OpenAI LLM 实现 - 支持工具调用与流式输出"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, AsyncGenerator, Dict, List, Optional
+
+from loom.interfaces.llm import BaseLLM
+
+try:
+    from openai import AsyncOpenAI
+except ImportError:
+    AsyncOpenAI = None  # type: ignore
+
+
+class OpenAILLM(BaseLLM):
+    """OpenAI LLM 实现 - 支持 gpt-4/gpt-3.5-turbo 等模型"""
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = "gpt-4",
+        base_url: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        if AsyncOpenAI is None:
+            raise ImportError("Please install openai package: pip install openai")
+
+        self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+        self._model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.kwargs = kwargs
+
+    @property
+    def model_name(self) -> str:
+        return self._model
+
+    @property
+    def supports_tools(self) -> bool:
+        # GPT-4 和 GPT-3.5-turbo 都支持工具调用
+        return "gpt-4" in self._model.lower() or "gpt-3.5-turbo" in self._model.lower()
+
+    async def generate(self, messages: List[Dict]) -> str:
+        """同步生成(非流式)"""
+        params = {
+            "model": self._model,
+            "messages": messages,
+            "temperature": self.temperature,
+        }
+        if self.max_tokens:
+            params["max_tokens"] = self.max_tokens
+
+        response = await self.client.chat.completions.create(**params)
+        return response.choices[0].message.content or ""
+
+    async def stream(self, messages: List[Dict]) -> AsyncGenerator[str, None]:
+        """流式生成"""
+        params = {
+            "model": self._model,
+            "messages": messages,
+            "temperature": self.temperature,
+            "stream": True,
+        }
+        if self.max_tokens:
+            params["max_tokens"] = self.max_tokens
+
+        stream = await self.client.chat.completions.create(**params)
+        async for chunk in stream:
+            content = chunk.choices[0].delta.content
+            if content:
+                yield content
+
+    async def generate_with_tools(self, messages: List[Dict], tools: List[Dict]) -> Dict:
+        """带工具调用的生成"""
+        params = {
+            "model": self._model,
+            "messages": messages,
+            "tools": tools,
+            "temperature": self.temperature,
+        }
+        if self.max_tokens:
+            params["max_tokens"] = self.max_tokens
+
+        response = await self.client.chat.completions.create(**params)
+        message = response.choices[0].message
+
+        # 解析工具调用
+        tool_calls = []
+        if message.tool_calls:
+            for tc in message.tool_calls:
+                tool_calls.append(
+                    {
+                        "id": tc.id,
+                        "name": tc.function.name,
+                        "arguments": json.loads(tc.function.arguments) if tc.function.arguments else {},
+                    }
+                )
+
+        return {
+            "content": message.content or "",
+            "tool_calls": tool_calls if tool_calls else None,
+        }
+
+    async def stream_with_tools(
+        self, messages: List[Dict], tools: List[Dict]
+    ) -> AsyncGenerator[Dict[str, Any], None]:
+        """流式工具调用(高级功能)"""
+        params = {
+            "model": self._model,
+            "messages": messages,
+            "tools": tools,
+            "temperature": self.temperature,
+            "stream": True,
+        }
+        if self.max_tokens:
+            params["max_tokens"] = self.max_tokens
+
+        stream = await self.client.chat.completions.create(**params)
+
+        # 累积工具调用信息
+        accumulated_tool_calls: Dict[int, Dict] = {}
+
+        async for chunk in stream:
+            delta = chunk.choices[0].delta
+
+            # 文本增量
+            if delta.content:
+                yield {"type": "text_delta", "content": delta.content}
+
+            # 工具调用增量
+            if delta.tool_calls:
+                for tc_delta in delta.tool_calls:
+                    idx = tc_delta.index
+                    if idx not in accumulated_tool_calls:
+                        accumulated_tool_calls[idx] = {
+                            "id": tc_delta.id or "",
+                            "name": "",
+                            "arguments": "",
+                        }
+
+                    if tc_delta.id:
+                        accumulated_tool_calls[idx]["id"] = tc_delta.id
+                    if tc_delta.function.name:
+                        accumulated_tool_calls[idx]["name"] = tc_delta.function.name
+                    if tc_delta.function.arguments:
+                        accumulated_tool_calls[idx]["arguments"] += tc_delta.function.arguments
+
+        # 流结束后,输出完整的工具调用
+        if accumulated_tool_calls:
+            tool_calls = []
+            for tc_data in accumulated_tool_calls.values():
+                try:
+                    arguments = json.loads(tc_data["arguments"]) if tc_data["arguments"] else {}
+                except json.JSONDecodeError:
+                    arguments = {}
+
+                tool_calls.append(
+                    {
+                        "id": tc_data["id"],
+                        "name": tc_data["name"],
+                        "arguments": arguments,
+                    }
+                )
+
+            yield {"type": "tool_calls", "tool_calls": tool_calls}

loom/builtin/llms/rule.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+import asyncio
+import re
+from typing import AsyncGenerator, Dict, List
+
+from loom.interfaces.llm import BaseLLM
+
+
+class RuleLLM(BaseLLM):
+    """规则驱动的演示 LLM：
+    - 若用户消息包含 "calc:" 前缀，则产生 calculator 的工具调用。
+    - 否则返回纯文本。
+    """
+
+    def __init__(self, model: str = "rule-llm") -> None:
+        self._model_name = model
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    @property
+    def supports_tools(self) -> bool:
+        return True
+
+    async def generate(self, messages: List[Dict]) -> str:
+        last = messages[-1]["content"] if messages else ""
+        return f"Echo: {last}"
+
+    async def stream(self, messages: List[Dict]) -> AsyncGenerator[str, None]:
+        text = await self.generate(messages)
+        for ch in text:
+            await asyncio.sleep(0)
+            yield ch
+
+    async def generate_with_tools(self, messages: List[Dict], tools: List[Dict]) -> Dict:
+        if not messages:
+            return {"content": "", "tool_calls": []}
+
+        last_msg = messages[-1]
+        last_role = last_msg.get("role")
+        last = last_msg.get("content", "")
+
+        # 若上条是工具结果，直接返回总结，不再调用工具
+        if last_role == "tool":
+            return {"content": f"工具执行完成，输出: {last}", "tool_calls": []}
+
+        # 简易规则：calc/read/write/glob/grep
+        m = re.search(r"calc:\s*([0-9+\-*/. ()]+)", last, re.IGNORECASE)
+        if m:
+            expr = m.group(1).strip()
+            return {
+                "content": "我将使用 calculator 计算该表达式。",
+                "tool_calls": [
+                    {"id": "tool_1", "name": "calculator", "arguments": {"expression": expr}}
+                ],
+            }
+
+        m = re.search(r"read:\s*(.+)$", last, re.IGNORECASE)
+        if m:
+            path = m.group(1).strip()
+            return {
+                "content": "我将读取指定文件。",
+                "tool_calls": [{"id": "tool_1", "name": "read_file", "arguments": {"path": path}}],
+            }
+
+        m = re.search(r"write:\s*([^<]+)<<<(.*)$", last, re.IGNORECASE | re.DOTALL)
+        if m:
+            path = m.group(1).strip()
+            content = m.group(2).strip()
+            return {
+                "content": "我将写入指定文件。",
+                "tool_calls": [
+                    {"id": "tool_1", "name": "write_file", "arguments": {"path": path, "content": content}}
+                ],
+            }
+
+        m = re.search(r"glob:\s*(.+)$", last, re.IGNORECASE)
+        if m:
+            pattern = m.group(1).strip()
+            return {
+                "content": "我将搜索匹配的文件。",
+                "tool_calls": [{"id": "tool_1", "name": "glob", "arguments": {"pattern": pattern}}],
+            }
+
+        m = re.search(r"grep:\s*(.+?)\s+in\s+(.+)$", last, re.IGNORECASE)
+        if m:
+            pattern = m.group(1).strip()
+            target = m.group(2).strip()
+            args = {"pattern": pattern}
+            if target.startswith("glob:"):
+                args["glob_pattern"] = target.split(":", 1)[1].strip()
+            else:
+                args["path"] = target
+            return {
+                "content": "我将检索文本匹配。",
+                "tool_calls": [{"id": "tool_1", "name": "grep", "arguments": args}],
+            }
+
+        # 默认：无工具
+        return {"content": await self.generate(messages), "tool_calls": []}

loom/builtin/memory/__init__.py

@@ -0,0 +1,5 @@
+from .in_memory import InMemoryMemory
+from .persistent_memory import PersistentMemory
+
+__all__ = ["InMemoryMemory", "PersistentMemory"]
+

loom/builtin/memory/in_memory.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from loom.core.types import Message
+from loom.interfaces.memory import BaseMemory
+
+
+class InMemoryMemory(BaseMemory):
+    def __init__(self) -> None:
+        self._messages: List[Message] = []
+
+    async def add_message(self, message: Message) -> None:
+        self._messages.append(message)
+
+    async def get_messages(self, limit: Optional[int] = None) -> List[Message]:
+        return self._messages[-limit:] if limit else list(self._messages)
+
+    async def clear(self) -> None:
+        self._messages.clear()
+