ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/memory/embeddings/__init__.py

@@ -0,0 +1 @@
+

chapter9/gaia-experience/AWorld/aworld/memory/embeddings/base.py

@@ -0,0 +1,97 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import List, Optional
+import uuid
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from aworld.core.memory import EmbeddingsConfig
+
+
+class EmbeddingsMetadata(BaseModel):
+    memory_id: str = Field(..., description="memory_id")
+    agent_id: Optional[str] = Field(default=None, description="agent_id")
+    session_id: Optional[str] = Field(default=None, description="session_id")
+    task_id: Optional[str] = Field(default=None, description="task_id")
+    user_id: Optional[str] = Field(default=None, description="user_id")
+    application_id: Optional[str] = Field(default=None, description="application_id")
+    memory_type: str = Field(..., description="memory_type")
+    embedding_model: str = Field(..., description="Embedding model")
+    created_at: str = Field(default_factory=lambda: datetime.now().isoformat(), description="Created at")
+    updated_at: str = Field(default_factory=lambda: datetime.now().isoformat(), description="Updated at")
+
+    model_config = ConfigDict(extra="allow")
+
+class EmbeddingsResult(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="ID")
+    embedding: Optional[list[float]] = Field(default=None, description="Embedding")
+    content: str = Field(..., description="Content")
+    metadata: Optional[EmbeddingsMetadata] = Field(..., description="Metadata")
+    score: Optional[float] = Field(default=None, description="Retrieved relevance score")
+
+class EmbeddingsResults(BaseModel):
+    docs: Optional[List[EmbeddingsResult]]
+    retrieved_at: int = Field(..., description="Retrieved at")
+
+class Embeddings(ABC):
+    """Interface for embedding models.
+    Embeddings are used to convert artifacts and queries into a vector space.
+    """
+    @abstractmethod
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        raise NotImplementedError
+
+    async def async_embed_query(self, text: str) -> list[float]:
+        """Asynchronous Embed query text."""
+        raise NotImplementedError
+
+
+class EmbeddingsBase(Embeddings):
+    """
+    Base class for embedding implementations that contains common functionality.
+    """
+
+    def __init__(self, config: EmbeddingsConfig):
+        """
+        Initialize EmbeddingsBase with configuration.
+        Args:
+            config (EmbeddingsConfig): Configuration for embedding model and API.
+        """
+        self.config = config
+
+
+    @abstractmethod
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Abstract method to embed a query string.
+        Args:
+            text (str): Text to embed.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        pass
+
+    @abstractmethod
+    async def async_embed_query(self, text: str) -> List[float]:
+        """
+        Abstract method to asynchronously embed a query string.
+        Args:
+            text (str): Text to embed.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        pass
+
+class EmbeddingFactory:
+
+    @staticmethod
+    def get_embedder(config: EmbeddingsConfig) -> Embeddings:
+        if config.provider == "openai":
+            from aworld.memory.embeddings.openai_compatible import OpenAICompatibleEmbeddings
+            return OpenAICompatibleEmbeddings(config)
+        elif config.provider == "ollama":
+            from aworld.memory.embeddings.ollama import OllamaEmbeddings
+            return OllamaEmbeddings(config)
+        else:
+            raise ValueError(f"Unsupported embedding provider: {config.provider}")

chapter9/gaia-experience/AWorld/aworld/memory/embeddings/factory.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+from aworld.core.memory import EmbeddingsConfig
+from aworld.memory.embeddings.base import Embeddings
+
+
+class EmbedderFactory:
+
+    @staticmethod
+    def get_embedder(config: EmbeddingsConfig) -> Optional[Embeddings]:
+        if not config:
+            return None
+        if config.provider == "openai":
+            from aworld.memory.embeddings.openai_compatible import OpenAICompatibleEmbeddings
+            return OpenAICompatibleEmbeddings(config)
+        elif config.provider == "ollama":
+            from aworld.memory.embeddings.ollama import OllamaEmbeddings
+            return OllamaEmbeddings(config)
+        else:
+            raise ValueError(f"Unsupported provider: {config.provider}")

chapter9/gaia-experience/AWorld/aworld/memory/embeddings/ollama.py

@@ -0,0 +1,82 @@
+import logging
+from typing import List
+
+import aiohttp
+import requests
+
+from aworld.core.memory import EmbeddingsConfig
+from aworld.memory.embeddings.base import EmbeddingsBase
+
+
+class OllamaEmbeddings(EmbeddingsBase):
+    """
+    Embedding implementation using Ollama HTTP API.
+    """
+    def __init__(self, config: EmbeddingsConfig):
+        """
+        Initialize OllamaEmbeddings with configuration.
+        Args:
+            config (EmbeddingsConfig): Configuration for embedding model and API.
+        """
+        super().__init__(config)
+    
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Embed a query string using Ollama HTTP API.
+        Args:
+            text (str): Text to embed.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        url = self.config.base_url.rstrip('/') + "/api/embed"
+        payload = {
+            "model": self.config.model_name,
+            "input": text
+        }
+        try:
+            response = requests.post(url, json=payload, timeout=self.config.timeout)
+            response.raise_for_status()
+            data = response.json()
+            # Ollama returns {"embedding": [...], ...}
+            logging.debug(f"Ollama embedding response: {data}")
+            return self.resolve_embedding(data)
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            raise RuntimeError(f"Ollama embedding API error: {e}")
+
+    async def async_embed_query(self, text: str) -> List[float]:
+        """
+        Asynchronously embed a query string using Ollama HTTP API.
+        Args:
+            text (str): Text to embed.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        url = self.config.base_url.rstrip('/') + "/api/embed"
+        payload = {
+            "model": self.config.model_name,
+            "input": text
+        }
+        try:
+            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.config.timeout)) as session:
+                async with session.post(url, json=payload) as resp:
+                    resp.raise_for_status()
+                    data = await resp.json()
+                    return self.resolve_embedding(data)
+        except Exception as e:
+            raise RuntimeError(f"Ollama async embedding API error: {e}")
+      
+    @staticmethod
+    def resolve_embedding(data: dict) -> List[float]:
+        """
+        Resolve the embedding from the response data.
+        Args:
+            data (dict): Response data from Ollama API.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        if "embeddings" in data and len(data["embeddings"]) > 0:
+            return data["embeddings"][0]
+        else:
+            return None

chapter9/gaia-experience/AWorld/aworld/memory/embeddings/openai_compatible.py

@@ -0,0 +1,79 @@
+import asyncio
+import time
+import logging
+from typing import Any, List
+
+from openai import OpenAI
+
+from aworld.core.memory import EmbeddingsConfig
+from aworld.memory.embeddings.base import EmbeddingsBase
+
+
+class OpenAICompatibleEmbeddings(EmbeddingsBase):
+    """
+    OpenAI compatible embeddings using OpenAI-compatible HTTP API.
+
+    - text-embedding-v4: [2048、1536、1024（默认）、768、512、256、128、64]
+    - text-embedding-v3: [1024(默认)、512、256、128、64]
+    - text-embedding-v2: [1536]
+    - text-embedding-v1: [1536]
+    """
+
+    def __init__(self, config: EmbeddingsConfig):
+        """
+        Initialize OpenAICompatibleEmbeddings with configuration.
+        Args:
+            config (EmbeddingsConfig): Configuration for embedding model and API.
+        """
+        super().__init__(config)
+        self.client = OpenAI(api_key=config.api_key, base_url=config.base_url)
+
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Embed a query string using OpenAI-compatible HTTP API.
+        Args:
+            text (str): Text to embed.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        try:
+            response = self.client.embeddings.create(
+                model=self.config.model_name,
+                input=text,
+                dimensions=self.config.dimensions)
+            data = response.data
+            logging.debug(f"OpenAI embedding response: {data}")
+            return self.resolve_embedding(data)
+        except Exception as e:
+            raise RuntimeError(f"OpenAI embedding API error: {e}")
+
+    async def async_embed_query(self, text: str) -> List[float]:
+        """
+        Asynchronously embed a query string using OpenAI-compatible HTTP API.
+        Args:
+            text (str): Text to embed.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        try:
+            response = self.client.embeddings.create(
+                model=self.config.model_name,
+                input=text,
+                dimensions=self.config.dimensions)
+            data = response.data
+            logging.debug(f"OpenAI embedding response: {data}")
+            return self.resolve_embedding(data)
+        except Exception as e:
+            raise RuntimeError(f"OpenAI async embedding API error: {e}")
+
+    @staticmethod
+    def resolve_embedding(data: list[Any]) -> List[float]:
+        """
+        Resolve the embedding from the response data (OpenAI format).
+        Args:
+            data (dict): Response data from OpenAI API.
+        Returns:
+            List[float]: Embedding vector.
+        """
+        return data[0].embedding