loom-agent 0.0.1__py3-none-any.whl

loom/builtin/retriever/qdrant_store.py

@@ -0,0 +1,274 @@
+"""Qdrant 向量存储适配器"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+import uuid
+
+from loom.interfaces.retriever import Document
+from loom.interfaces.vector_store import BaseVectorStore
+
+try:
+    from qdrant_client import QdrantClient
+    from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
+    QDRANT_AVAILABLE = True
+except ImportError:
+    QDRANT_AVAILABLE = False
+
+
+class QdrantVectorStore(BaseVectorStore):
+    """
+    Qdrant 向量存储适配器
+
+    特点:
+    - ✅ 开源向量数据库
+    - ✅ 支持本地部署和云服务
+    - ✅ 高性能 Rust 实现
+    - ✅ 丰富的过滤功能
+    - ✅ 支持 gRPC 和 HTTP
+
+    示例:
+        from loom.builtin.retriever.qdrant_store import QdrantVectorStore
+        from loom.builtin.retriever.vector_store_config import QdrantConfig
+
+        # 本地 Qdrant
+        config = QdrantConfig.create(
+            host="localhost",
+            port=6333,
+            collection_name="loom_docs"
+        )
+
+        # Qdrant Cloud
+        config = QdrantConfig.create(
+            host="your-cluster.qdrant.io",
+            api_key="your-api-key",
+            https=True
+        )
+
+        vector_store = QdrantVectorStore(config)
+        await vector_store.initialize()
+    """
+
+    def __init__(self, config: Dict[str, Any] | Any):
+        """
+        Parameters:
+            config: QdrantConfig 对象或配置字典
+        """
+        if not QDRANT_AVAILABLE:
+            raise ImportError(
+                "Qdrant client is not installed. "
+                "Install with: pip install qdrant-client"
+            )
+
+        # 支持字典或 Pydantic 模型
+        if hasattr(config, "model_dump"):
+            self.config = config.model_dump()
+        else:
+            self.config = config
+
+        self.host = self.config.get("host", "localhost")
+        self.port = self.config.get("port", 6333)
+        self.collection_name = self.config.get("collection_name", "loom_documents")
+        self.dimension = self.config.get("dimension", 1536)
+        self.metric = self.config.get("metric", "cosine")
+        self.api_key = self.config.get("api_key")
+        self.https = self.config.get("https", False)
+        self.prefer_grpc = self.config.get("prefer_grpc", False)
+
+        self.client: Optional[QdrantClient] = None
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """初始化 Qdrant 连接和集合"""
+        if self._initialized:
+            return
+
+        # 初始化客户端
+        self.client = QdrantClient(
+            host=self.host,
+            port=self.port,
+            api_key=self.api_key,
+            https=self.https,
+            prefer_grpc=self.prefer_grpc,
+        )
+
+        # 检查集合是否存在
+        collections = self.client.get_collections().collections
+        collection_names = [c.name for c in collections]
+
+        if self.collection_name not in collection_names:
+            # 创建集合
+            distance_map = {
+                "cosine": Distance.COSINE,
+                "euclidean": Distance.EUCLID,
+                "dot_product": Distance.DOT,
+            }
+
+            self.client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=VectorParams(
+                    size=self.dimension,
+                    distance=distance_map.get(self.metric, Distance.COSINE)
+                )
+            )
+
+        self._initialized = True
+
+    async def add_vectors(
+        self,
+        vectors: List[List[float]],
+        documents: List[Document]
+    ) -> None:
+        """
+        添加向量到 Qdrant
+
+        Parameters:
+            vectors: 向量列表
+            documents: 对应的文档列表
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        # 构建 Qdrant Point 格式
+        points = []
+        for i, (vector, doc) in enumerate(zip(vectors, documents)):
+            # 生成或使用文档 ID
+            point_id = doc.doc_id or str(uuid.uuid4())
+
+            # 构建 payload（元数据 + 内容）
+            payload = {
+                "content": doc.content,
+                **(doc.metadata or {})
+            }
+
+            # 添加分数（如果有）
+            if doc.score is not None:
+                payload["score"] = doc.score
+
+            points.append(
+                PointStruct(
+                    id=point_id,
+                    vector=vector,
+                    payload=payload
+                )
+            )
+
+        # 批量上传
+        self.client.upsert(
+            collection_name=self.collection_name,
+            points=points
+        )
+
+    async def search(
+        self,
+        query_vector: List[float],
+        top_k: int = 5,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> List[Tuple[Document, float]]:
+        """
+        搜索相似向量
+
+        Parameters:
+            query_vector: 查询向量
+            top_k: 返回结果数量
+            filters: 元数据过滤条件
+
+        Returns:
+            [(Document, score), ...] 列表
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        # 构建过滤器
+        qdrant_filter = None
+        if filters:
+            qdrant_filter = self._build_qdrant_filter(filters)
+
+        # 执行查询
+        results = self.client.search(
+            collection_name=self.collection_name,
+            query_vector=query_vector,
+            limit=top_k,
+            query_filter=qdrant_filter
+        )
+
+        # 转换结果
+        documents_with_scores = []
+        for hit in results:
+            # 提取内容和元数据
+            payload = hit.payload
+            content = payload.pop("content", "")
+            score = payload.pop("score", None)
+
+            doc = Document(
+                content=content,
+                metadata=payload,
+                score=hit.score,
+                doc_id=str(hit.id)
+            )
+            documents_with_scores.append((doc, hit.score))
+
+        return documents_with_scores
+
+    async def delete(self, doc_ids: List[str]) -> None:
+        """
+        删除文档
+
+        Parameters:
+            doc_ids: 文档 ID 列表
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        self.client.delete(
+            collection_name=self.collection_name,
+            points_selector=doc_ids
+        )
+
+    async def clear(self) -> None:
+        """清空集合"""
+        if not self._initialized:
+            await self.initialize()
+
+        # 删除并重建集合
+        self.client.delete_collection(collection_name=self.collection_name)
+        await self.initialize()
+
+    def _build_qdrant_filter(self, filters: Dict[str, Any]) -> Filter:
+        """
+        构建 Qdrant 过滤器
+
+        示例:
+            {"category": "python", "price": {"$gte": 100}}
+            →
+            Filter(
+                must=[
+                    FieldCondition(key="category", match=MatchValue(value="python")),
+                    FieldCondition(key="price", range=Range(gte=100))
+                ]
+            )
+        """
+        conditions = []
+
+        for key, value in filters.items():
+            if isinstance(value, dict):
+                # 复杂查询（范围/比较）
+                # 简化实现：仅支持基本相等匹配
+                # 生产环境可扩展支持 $gte, $lte, $in 等
+                pass
+            else:
+                # 简单相等查询
+                conditions.append(
+                    FieldCondition(
+                        key=key,
+                        match=MatchValue(value=value)
+                    )
+                )
+
+        return Filter(must=conditions) if conditions else None
+
+    async def close(self) -> None:
+        """关闭连接"""
+        if self.client:
+            self.client.close()
+        self._initialized = False

loom/builtin/retriever/vector_store.py

@@ -0,0 +1,128 @@
+"""基于向量存储的检索器实现"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from loom.interfaces.retriever import BaseRetriever, Document
+from loom.interfaces.vector_store import BaseVectorStore
+from loom.interfaces.embedding import BaseEmbedding
+
+
+class VectorStoreRetriever(BaseRetriever):
+    """
+    向量存储检索器 - 基于语义相似度检索
+
+    将向量存储和 Embedding 模型组合，提供语义检索能力。
+
+    特点:
+    - ✅ 语义相似度检索
+    - ✅ 支持多种向量数据库
+    - ✅ 支持多种 Embedding 模型
+    - ✅ 自动向量化
+
+    示例:
+        from loom.builtin.retriever.vector_store import VectorStoreRetriever
+        from loom.builtin.retriever.chroma_store import ChromaVectorStore
+        from loom.builtin.embeddings import OpenAIEmbedding
+
+        vector_store = ChromaVectorStore(config)
+        await vector_store.initialize()
+
+        embedding = OpenAIEmbedding(api_key="...")
+
+        retriever = VectorStoreRetriever(
+            vector_store=vector_store,
+            embedding=embedding
+        )
+
+        # 添加文档（自动向量化）
+        await retriever.add_documents([
+            Document(content="Loom is an AI agent framework"),
+            Document(content="Loom supports RAG capabilities"),
+        ])
+
+        # 检索（自动向量化查询）
+        results = await retriever.retrieve("What is Loom?", top_k=2)
+    """
+
+    def __init__(
+        self,
+        vector_store: BaseVectorStore,
+        embedding: BaseEmbedding,
+    ):
+        """
+        Parameters:
+            vector_store: 向量存储实例（Pinecone, Qdrant, Milvus, ChromaDB）
+            embedding: Embedding 模型实例（OpenAI, Sentence Transformers）
+        """
+        self.vector_store = vector_store
+        self.embedding = embedding
+
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int = 5,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> List[Document]:
+        """
+        基于语义相似度检索文档
+
+        Parameters:
+            query: 查询文本
+            top_k: 返回文档数量
+            filters: 元数据过滤条件（可选）
+
+        Returns:
+            Document 列表，按相似度分数降序排列
+        """
+        # Step 1: 将查询向量化
+        query_vector = await self.embedding.embed_query(query)
+
+        # Step 2: 在向量存储中搜索
+        results = await self.vector_store.search(
+            query_vector=query_vector,
+            top_k=top_k,
+            filters=filters
+        )
+
+        # Step 3: 提取文档（已包含 score）
+        documents = [doc for doc, score in results]
+
+        return documents
+
+    async def add_documents(self, documents: List[Document]) -> None:
+        """
+        添加文档到向量存储（自动向量化）
+
+        Parameters:
+            documents: 文档列表
+        """
+        if not documents:
+            return
+
+        # Step 1: 提取文档内容
+        texts = [doc.content for doc in documents]
+
+        # Step 2: 批量向量化
+        vectors = await self.embedding.embed_documents(texts)
+
+        # Step 3: 添加到向量存储
+        await self.vector_store.add_vectors(vectors, documents)
+
+    async def delete_documents(self, doc_ids: List[str]) -> None:
+        """
+        删除文档
+
+        Parameters:
+            doc_ids: 文档 ID 列表
+        """
+        await self.vector_store.delete(doc_ids)
+
+    async def clear(self) -> None:
+        """清空所有文档"""
+        await self.vector_store.clear()
+
+    async def close(self) -> None:
+        """关闭连接"""
+        await self.vector_store.close()

loom/builtin/retriever/vector_store_config.py

@@ -0,0 +1,217 @@
+"""向量数据库配置管理 - 统一的配置接口"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+from pydantic import BaseModel, Field
+
+
+class VectorStoreConfig(BaseModel):
+    """向量数据库配置基类"""
+    store_type: str = Field(description="向量数据库类型: pinecone, qdrant, milvus, chroma")
+    api_key: Optional[str] = Field(default=None, description="API Key（如果需要）")
+    host: Optional[str] = Field(default=None, description="数据库地址")
+    port: Optional[int] = Field(default=None, description="数据库端口")
+    collection_name: str = Field(default="loom_documents", description="集合/索引名称")
+    dimension: int = Field(default=1536, description="向量维度")
+    metric: str = Field(default="cosine", description="相似度度量: cosine, euclidean, dot_product")
+    extra_params: Dict[str, Any] = Field(default_factory=dict, description="额外参数")
+
+
+class PineconeConfig(VectorStoreConfig):
+    """Pinecone 配置"""
+    store_type: str = "pinecone"
+    environment: str = Field(description="Pinecone 环境: us-west1-gcp, eu-west1-gcp, etc.")
+    index_name: str = Field(description="索引名称")
+    namespace: Optional[str] = Field(default=None, description="命名空间（可选）")
+
+    @classmethod
+    def create(
+        cls,
+        api_key: str,
+        environment: str,
+        index_name: str,
+        namespace: Optional[str] = None,
+        dimension: int = 1536,
+    ) -> "PineconeConfig":
+        """快速创建 Pinecone 配置"""
+        return cls(
+            api_key=api_key,
+            environment=environment,
+            index_name=index_name,
+            namespace=namespace,
+            dimension=dimension,
+        )
+
+
+class QdrantConfig(VectorStoreConfig):
+    """Qdrant 配置"""
+    store_type: str = "qdrant"
+    host: str = Field(default="localhost", description="Qdrant 服务地址")
+    port: int = Field(default=6333, description="Qdrant 端口")
+    grpc_port: Optional[int] = Field(default=None, description="gRPC 端口（可选）")
+    prefer_grpc: bool = Field(default=False, description="优先使用 gRPC")
+    https: bool = Field(default=False, description="是否使用 HTTPS")
+
+    @classmethod
+    def create(
+        cls,
+        host: str = "localhost",
+        port: int = 6333,
+        collection_name: str = "loom_documents",
+        dimension: int = 1536,
+        api_key: Optional[str] = None,
+        https: bool = False,
+    ) -> "QdrantConfig":
+        """快速创建 Qdrant 配置"""
+        return cls(
+            host=host,
+            port=port,
+            collection_name=collection_name,
+            dimension=dimension,
+            api_key=api_key,
+            https=https,
+        )
+
+
+class MilvusConfig(VectorStoreConfig):
+    """Milvus 配置"""
+    store_type: str = "milvus"
+    host: str = Field(default="localhost", description="Milvus 服务地址")
+    port: int = Field(default=19530, description="Milvus 端口")
+    user: Optional[str] = Field(default=None, description="用户名")
+    password: Optional[str] = Field(default=None, description="密码")
+    secure: bool = Field(default=False, description="是否使用安全连接")
+    index_type: str = Field(default="IVF_FLAT", description="索引类型: IVF_FLAT, HNSW, etc.")
+    index_params: Dict[str, Any] = Field(default_factory=dict, description="索引参数")
+
+    @classmethod
+    def create(
+        cls,
+        host: str = "localhost",
+        port: int = 19530,
+        collection_name: str = "loom_documents",
+        dimension: int = 1536,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        index_type: str = "IVF_FLAT",
+    ) -> "MilvusConfig":
+        """快速创建 Milvus 配置"""
+        return cls(
+            host=host,
+            port=port,
+            collection_name=collection_name,
+            dimension=dimension,
+            user=user,
+            password=password,
+            index_type=index_type,
+        )
+
+
+class ChromaConfig(VectorStoreConfig):
+    """ChromaDB 配置"""
+    store_type: str = "chroma"
+    host: Optional[str] = Field(default=None, description="ChromaDB 服务地址（远程模式）")
+    port: Optional[int] = Field(default=8000, description="ChromaDB 端口（远程模式）")
+    persist_directory: Optional[str] = Field(default=None, description="持久化目录（本地模式）")
+    client_type: str = Field(default="local", description="客户端类型: local, http")
+
+    @classmethod
+    def create_local(
+        cls,
+        persist_directory: str = "./chroma_db",
+        collection_name: str = "loom_documents",
+        dimension: int = 1536,
+    ) -> "ChromaConfig":
+        """创建本地 ChromaDB 配置"""
+        return cls(
+            persist_directory=persist_directory,
+            collection_name=collection_name,
+            dimension=dimension,
+            client_type="local",
+        )
+
+    @classmethod
+    def create_remote(
+        cls,
+        host: str = "localhost",
+        port: int = 8000,
+        collection_name: str = "loom_documents",
+        dimension: int = 1536,
+    ) -> "ChromaConfig":
+        """创建远程 ChromaDB 配置"""
+        return cls(
+            host=host,
+            port=port,
+            collection_name=collection_name,
+            dimension=dimension,
+            client_type="http",
+        )
+
+
+class EmbeddingConfig(BaseModel):
+    """Embedding 模型配置"""
+    provider: str = Field(description="提供商: openai, huggingface, cohere, sentence_transformers")
+    model_name: str = Field(description="模型名称")
+    api_key: Optional[str] = Field(default=None, description="API Key（如果需要）")
+    dimension: int = Field(default=1536, description="向量维度")
+    batch_size: int = Field(default=32, description="批处理大小")
+    extra_params: Dict[str, Any] = Field(default_factory=dict, description="额外参数")
+
+    @classmethod
+    def openai(
+        cls,
+        api_key: str,
+        model_name: str = "text-embedding-3-small",
+        dimension: int = 1536,
+    ) -> "EmbeddingConfig":
+        """OpenAI Embedding 配置"""
+        return cls(
+            provider="openai",
+            model_name=model_name,
+            api_key=api_key,
+            dimension=dimension,
+        )
+
+    @classmethod
+    def huggingface(
+        cls,
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        dimension: int = 384,
+        api_key: Optional[str] = None,
+    ) -> "EmbeddingConfig":
+        """HuggingFace Embedding 配置"""
+        return cls(
+            provider="huggingface",
+            model_name=model_name,
+            api_key=api_key,
+            dimension=dimension,
+        )
+
+    @classmethod
+    def sentence_transformers(
+        cls,
+        model_name: str = "all-MiniLM-L6-v2",
+        dimension: int = 384,
+    ) -> "EmbeddingConfig":
+        """Sentence Transformers 本地模型配置"""
+        return cls(
+            provider="sentence_transformers",
+            model_name=model_name,
+            dimension=dimension,
+        )
+
+    @classmethod
+    def cohere(
+        cls,
+        api_key: str,
+        model_name: str = "embed-english-v3.0",
+        dimension: int = 1024,
+    ) -> "EmbeddingConfig":
+        """Cohere Embedding 配置"""
+        return cls(
+            provider="cohere",
+            model_name=model_name,
+            api_key=api_key,
+            dimension=dimension,
+        )

loom/builtin/tools/__init__.py

@@ -0,0 +1,32 @@
+from .calculator import Calculator
+from .read_file import ReadFileTool
+from .write_file import WriteFileTool
+from .glob import GlobTool
+from .grep import GrepTool
+
+try:
+    from .web_search import WebSearchTool
+    _has_web_search = True
+except ImportError:
+    _has_web_search = False
+
+try:
+    from .python_repl import PythonREPLTool
+    _has_python_repl = True
+except ImportError:
+    _has_python_repl = False
+
+try:
+    from .http_request import HTTPRequestTool
+    _has_http_request = True
+except ImportError:
+    _has_http_request = False
+
+__all__ = ["Calculator", "ReadFileTool", "WriteFileTool", "GlobTool", "GrepTool"]
+
+if _has_web_search:
+    __all__.append("WebSearchTool")
+if _has_python_repl:
+    __all__.append("PythonREPLTool")
+if _has_http_request:
+    __all__.append("HTTPRequestTool")

loom/builtin/tools/calculator.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import ast
+import operator as op
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from loom.interfaces.tool import BaseTool
+
+
+class CalcArgs(BaseModel):
+    expression: str = Field(description="Arithmetic expression, e.g., '2+2*3'")
+
+
+class Calculator(BaseTool):
+    name = "calculator"
+    description = "Evaluate simple arithmetic expressions"
+    args_schema = CalcArgs
+
+    async def run(self, **kwargs) -> Any:
+        expr = kwargs.get("expression", "")
+        return str(_safe_eval(expr))
+
+
+# 安全 eval：仅支持基本算术
+_ops = {
+    ast.Add: op.add,
+    ast.Sub: op.sub,
+    ast.Mult: op.mul,
+    ast.Div: op.truediv,
+    ast.Pow: op.pow,
+    ast.USub: op.neg,
+}
+
+
+def _safe_eval(expr: str) -> Any:
+    def _eval(node: ast.AST) -> Any:
+        if isinstance(node, ast.Num):
+            return node.n
+        if isinstance(node, ast.BinOp):
+            return _ops[type(node.op)](_eval(node.left), _eval(node.right))
+        if isinstance(node, ast.UnaryOp):
+            return _ops[type(node.op)](_eval(node.operand))
+        raise ValueError("Unsupported expression")
+
+    tree = ast.parse(expr, mode="eval")
+    return _eval(tree.body)
+