pycityagent 2.0.0a22__py3-none-any.whl → 2.0.0a25__py3-none-any.whl

pycityagent/__init__.py

@@ -5,6 +5,7 @@ Pycityagent: 城市智能体构建框架
 from .agent import Agent, CitizenAgent, InstitutionAgent
 from .environment import Simulator
 import logging
+from .llm import SentenceEmbedding
 
 # 创建一个 pycityagent 记录器
 logger = logging.getLogger("pycityagent")
@@ -19,4 +20,4 @@ if not logger.hasHandlers():
     handler.setFormatter(formatter)
     logger.addHandler(handler)
 
-__all__ = ["Agent", "Simulator", "CitizenAgent", "InstitutionAgent"]
+__all__ = ["Agent", "Simulator", "CitizenAgent", "InstitutionAgent","SentenceEmbedding",]

pycityagent/agent.py

@@ -236,7 +236,15 @@ class Agent(ABC):
 
         # 添加记忆上下文
         if self._memory:
-            relevant_memories = await self._memory.search(survey_prompt)
+            relevant_memories = await self.memory.search(survey_prompt)
+
+            formatted_results = []
+            # for result in top_results:
+            #     formatted_results.append(
+            #         f"- [{result['type']}] {result['content']} "
+            #         f"(相关度: {result['similarity']:.2f})"
+            #     )
+
             if relevant_memories:
                 dialog.append(
                     {
@@ -458,7 +466,9 @@ class Agent(ABC):
         topic = f"exps/{self._exp_id}/agents/{to_agent_uuid}/{sub_topic}"
         await self._messager.send_message(topic, payload)
 
-    async def send_message_to_agent(self, to_agent_uuid: str, content: str, type: str = "social"):
+    async def send_message_to_agent(
+        self, to_agent_uuid: str, content: str, type: str = "social"
+    ):
         """通过 Messager 发送消息"""
         if self._messager is None:
             raise RuntimeError("Messager is not set")
@@ -598,6 +608,7 @@ class CitizenAgent(Agent):
                 # 防止模拟器还没有到prepare阶段导致get_person出错
             self._has_bound_to_simulator = True
             self._agent_id = person_id
+            self.memory.set_agent_id(person_id)
 
     async def _bind_to_economy(self):
         if self._economy_client is None:

pycityagent/environment/simulator.py

@@ -3,7 +3,6 @@
 import asyncio
 import logging
 import os
-from collections.abc import Sequence
 from datetime import datetime, timedelta
 from typing import Any, Optional, Union, cast
 
@@ -22,13 +21,14 @@ from .utils.const import *
 
 logger = logging.getLogger("pycityagent")
 
+
 class Simulator:
     """
     - 模拟器主类
     - Simulator Class
     """
 
-    def __init__(self, config:dict, secure: bool = False) -> None:
+    def __init__(self, config: dict, secure: bool = False) -> None:
         self.config = config
         """
         - 模拟器配置
@@ -193,7 +193,7 @@ class Simulator:
         else:
             # BUG: 返回的time是float类型
             return t_sec["t"]
-        
+
     async def get_simulator_day(self) -> int:
         """
         获取模拟器到第几日
@@ -202,7 +202,7 @@ class Simulator:
         t_sec = cast(dict[str, int], t_sec)
         day = t_sec["t"] // 86400
         return day
-    
+
     async def get_simulator_second_from_start_of_day(self) -> int:
         """
         获取模拟器从00:00:00到当前的秒数
@@ -316,7 +316,7 @@ class Simulator:
         radius: float,
         poi_type: Union[str, list[str]],
     ):
-        if not isinstance(poi_type, Sequence):
+        if isinstance(poi_type, str):
             poi_type = [poi_type]
         transformed_poi_type = []
         for t in poi_type:

pycityagent/llm/__init__.py

@@ -1,6 +1,11 @@
 """LLM相关模块"""
 
+from .embeddings import SentenceEmbedding, SimpleEmbedding
 from .llm import LLM, LLMConfig
-from .embedding import SimpleEmbedding
 
-__all__ = ["LLM", "LLMConfig", "SimpleEmbedding"]
+__all__ = [
+    "LLM",
+    "LLMConfig",
+    "SentenceEmbedding",
+    "SimpleEmbedding",
+]

pycityagent/llm/embeddings.py

@@ -0,0 +1,231 @@
+import hashlib
+import json
+import os
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from langchain_core.embeddings import Embeddings
+from transformers import AutoModel, AutoTokenizer
+
+__all__ = [
+    "SentenceEmbedding",
+    "SimpleEmbedding",
+]
+
+
+class SentenceEmbedding(Embeddings):
+    def __init__(
+        self,
+        pretrained_model_name_or_path: Union[str, os.PathLike] = "BAAI/bge-m3",
+        max_seq_len: int = 8192,
+        auto_cuda: bool = False,
+        local_files_only: bool = False,
+        cache_dir: str = "./cache",
+        proxies: Optional[dict] = None,
+    ):
+        os.makedirs(cache_dir, exist_ok=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            proxies=proxies,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        self.model = AutoModel.from_pretrained(
+            pretrained_model_name_or_path,
+            proxies=proxies,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        self._cuda = auto_cuda and torch.cuda.is_available()
+
+        if self._cuda:
+            self.model = self.model.cuda()
+
+        self.model.eval()
+        self.max_seq_len = max_seq_len
+
+    def _embed(self, texts: list[str]) -> list[list[float]]:
+        # Tokenize sentences
+        encoded_input = self.tokenizer(
+            texts, padding=True, truncation=True, return_tensors="pt"
+        )
+        # for s2p(short query to long passage) retrieval task, add an instruction to query (not add instruction for passages)
+        # encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')
+
+        # check length of input
+        # assert seq_len <= 8192
+        assert encoded_input["input_ids"].shape[1] <= self.max_seq_len  # type: ignore
+
+        if self._cuda:
+            encoded_input = {k: v.cuda() for k, v in encoded_input.items()}
+        # Compute token embeddings
+        with torch.no_grad():
+            model_output = self.model(**encoded_input)
+            # Perform pooling. In this case, cls pooling.
+            sentence_embeddings = model_output[0][:, 0]
+        # normalize embeddings
+        sentence_embeddings = torch.nn.functional.normalize(
+            sentence_embeddings, p=2, dim=1
+        )
+        if self._cuda:
+            sentence_embeddings = sentence_embeddings.cpu()
+        return sentence_embeddings.tolist()
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents."""
+        return self._embed(texts)
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        return self._embed([text])[0]
+
+
+class SimpleEmbedding(Embeddings):
+    """简单的基于内存的embedding实现
+
+    使用简单的词袋模型(Bag of Words)和TF-IDF来生成文本的向量表示。
+    所有向量都保存在内存中，适用于小规模应用。
+    """
+
+    def __init__(self, vector_dim: int = 128, cache_size: int = 1000):
+        """初始化
+
+        Args:
+            vector_dim: 向量维度
+            cache_size: 缓存大小，超过此大小将清除最早的缓存
+        """
+        self.vector_dim = vector_dim
+        self.cache_size = cache_size
+        self._cache: dict[str, list[float]] = {}
+        self._vocab: dict[str, int] = {}  # 词汇表
+        self._idf: dict[str, float] = {}  # 逆文档频率
+        self._doc_count = 0  # 文档总数
+
+    def _text_to_hash(self, text: str) -> str:
+        """将文本转换为hash值"""
+        return hashlib.md5(text.encode()).hexdigest()
+
+    def _tokenize(self, text: str) -> list[str]:
+        """简单的分词"""
+        # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
+        return text.lower().split()
+
+    def _update_vocab(self, tokens: list[str]):
+        """更新词汇表"""
+        for token in set(tokens):  # 使用set去重
+            if token not in self._vocab:
+                self._vocab[token] = len(self._vocab)
+
+    def _update_idf(self, tokens: list[str]):
+        """更新IDF值"""
+        self._doc_count += 1
+        unique_tokens = set(tokens)
+        for token in unique_tokens:
+            self._idf[token] = self._idf.get(token, 0) + 1
+
+    def _calculate_tf(self, tokens: list[str]) -> dict[str, float]:
+        """计算词频(TF)"""
+        tf = {}
+        total_tokens = len(tokens)
+        for token in tokens:
+            tf[token] = tf.get(token, 0) + 1
+        # 归一化
+        for token in tf:
+            tf[token] /= total_tokens
+        return tf
+
+    def _calculate_tfidf(self, tokens: list[str]) -> list[float]:
+        """计算TF-IDF向量"""
+        vector = np.zeros(self.vector_dim)
+        tf = self._calculate_tf(tokens)
+
+        for token, tf_value in tf.items():
+            if token in self._idf:
+                idf = np.log(self._doc_count / self._idf[token])
+                idx = self._vocab[token] % self.vector_dim  # 使用取模运算来控制向量维度
+                vector[idx] += tf_value * idf
+
+        # L2归一化
+        norm = np.linalg.norm(vector)
+        if norm > 0:
+            vector /= norm
+
+        return list(vector)
+
+    def _embed(self, text: str) -> list[float]:
+        """生成文本的向量表示
+
+        Args:
+            text: 输入文本
+
+        Returns:
+            np.ndarray: 文本的向量表示
+        """
+        # 检查缓存
+        text_hash = self._text_to_hash(text)
+        if text_hash in self._cache:
+            return self._cache[text_hash]
+
+        # 分词
+        tokens = self._tokenize(text)
+        if not tokens:
+            return list(np.zeros(self.vector_dim))
+
+        # 更新词汇表和IDF
+        self._update_vocab(tokens)
+        self._update_idf(tokens)
+
+        # 计算向量
+        vector = self._calculate_tfidf(tokens)
+
+        # 更新缓存
+        if len(self._cache) >= self.cache_size:
+            # 删除最早的缓存
+            oldest_key = next(iter(self._cache))
+            del self._cache[oldest_key]
+        self._cache[text_hash] = vector
+
+        return list(vector)
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents."""
+        return [self._embed(text) for text in texts]
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        return self._embed(text)
+
+    # def save(self, file_path: str):
+    #     """保存模型"""
+    #     state = {
+    #         "vector_dim": self.vector_dim,
+    #         "cache_size": self.cache_size,
+    #         "vocab": self._vocab,
+    #         "idf": self._idf,
+    #         "doc_count": self._doc_count,
+    #     }
+    #     with open(file_path, "w") as f:
+    #         json.dump(state, f)
+
+    # def load(self, file_path: str):
+    #     """加载模型"""
+    #     with open(file_path, "r") as f:
+    #         state = json.load(f)
+    #     self.vector_dim = state["vector_dim"]
+    #     self.cache_size = state["cache_size"]
+    #     self._vocab = state["vocab"]
+    #     self._idf = state["idf"]
+    #     self._doc_count = state["doc_count"]
+    #     self._cache = {}  # 清空缓存
+
+
+if __name__ == "__main__":
+    # se = SentenceEmbedding(
+    #     pretrained_model_name_or_path="ignore/BAAI--bge-m3", cache_dir="ignore"
+    # )
+    se = SimpleEmbedding()
+    print(se.embed_query("hello world"))
+    print(se.embed_query("hello world"))
+    print(se.embed_query("hello world"))
+    print(se.embed_query("hello world"))

pycityagent/memory/__init__.py

@@ -1,5 +1,6 @@
 """Memory."""
 
+from .faiss_query import FaissQuery
 from .memory import Memory
 from .memory_base import MemoryBase, MemoryUnit
 from .profile import ProfileMemory, ProfileMemoryUnit
@@ -8,4 +9,5 @@ from .state import StateMemory
 
 __all__ = [
     "Memory",
+    "FaissQuery",
 ]

pycityagent/memory/faiss_query.py

@@ -0,0 +1,302 @@
+import asyncio
+from collections.abc import Sequence
+from typing import Any, Literal, Optional, Union
+
+import faiss
+import numpy as np
+from langchain_community.docstore.in_memory import InMemoryDocstore
+from langchain_community.vectorstores import FAISS
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+
+from ..utils.decorators import lock_decorator
+
+
+class FaissQuery:
+    def __init__(
+        self,
+        embeddings: Optional[Embeddings] = None,
+        index_type: Any = faiss.IndexFlatL2,
+        dimension: Optional[int] = None,
+    ) -> None:
+        self._embeddings = embeddings
+        self._lock = asyncio.Lock()
+        if embeddings is None:
+            self._index = None
+            self._vectors_store = None
+        else:
+            if dimension is None:
+                dimension = len(embeddings.embed_query("hello world"))
+            self._index = index_type(dimension)
+            self._vectors_store = FAISS(
+                embedding_function=embeddings,
+                index=self._index,
+                docstore=InMemoryDocstore(),
+                index_to_docstore_id={},
+            )
+
+    @property
+    def embeddings(
+        self,
+    ) -> Embeddings:
+        if self._embeddings is None:
+            raise RuntimeError(f"No embedding set, please `set_embeddings` first!")
+        return self._embeddings
+
+    @property
+    def vectors_store(
+        self,
+    ) -> FAISS:
+        if self._vectors_store is None:
+            raise RuntimeError(f"No embedding set, thus no vector stores initialized!")
+        return self._vectors_store
+
+    @lock_decorator
+    async def add_documents(
+        self,
+        agent_id: int,
+        documents: Union[str, Sequence[str]],
+        extra_tags: Optional[dict] = None,
+    ) -> list[str]:
+        if isinstance(documents, str):
+            documents = [documents]
+        _metadata = {"_id": agent_id}
+        if extra_tags is not None:
+            _metadata.update(extra_tags)
+        to_add_documents = [
+            Document(page_content=doc, metadata=_metadata) for doc in documents
+        ]
+        return await self.vectors_store.aadd_documents(
+            documents=to_add_documents,
+        )
+
+    @lock_decorator
+    async def delete_documents(
+        self,
+        to_delete_ids: list[str],
+    ):
+        await self.vectors_store.adelete(
+            ids=to_delete_ids,
+        )
+
+    @lock_decorator
+    async def similarity_search(
+        self,
+        query: str,
+        agent_id: int,
+        k: int = 4,
+        fetch_k: int = 20,
+        return_score_type: Union[
+            Literal["none"], Literal["similarity_score"], Literal["L2-distance"]
+        ] = "none",
+        filter: Optional[dict] = None,
+    ) -> Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]:
+        """
+        Return content most similar to the given query.
+
+        Args:
+            query (str): The text to look up documents similar to.
+            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
+            k (int, optional): The number of top similar contents to return. Defaults to 4.
+            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            return_score_type (Union[Literal["none"], Literal["similarity_score"], Literal["L2-distance"]], optional):
+                Specifies whether and how to return similarity scores with the results:
+                    - "none": Do not return scores; only return the contents (default).
+                    - "similarity_score": Return a tuple of content and its similarity score.
+                    - "L2-distance": Return a tuple of content and its L2 distance from the query.
+            filter (dict, optional): The filter dict for metadata.
+
+        Returns:
+            Union[list[tuple[str,dict]], list[tuple[str, float,dict]]]:
+                Depending on the `return_score_type` parameter, returns either a list of strings representing the top-k similar contents,
+                or a list of tuples where each tuple contains a string and a floating-point score.
+        """
+        _filter = {
+            "_id": agent_id,
+        }
+        if filter is not None:
+            _filter.update(filter)
+        if return_score_type == "L2-distance":
+            _result = await self.vectors_store.asimilarity_search_with_score(
+                query=query,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+            )
+            return [(r.page_content, s, r.metadata) for r, s in _result]
+        elif return_score_type == "none":
+            _result = await self.vectors_store.asimilarity_search(
+                query=query,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+            )
+            return [(r.page_content, r.metadata) for r in _result]
+        elif return_score_type == "similarity_score":
+            _result = await self.vectors_store.asimilarity_search_with_relevance_scores(
+                query=query,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+            )
+            return [(r.page_content, s, r.metadata) for r, s in _result]
+        else:
+            raise ValueError(f"Invalid `return_score_type` {return_score_type}!")
+
+    @lock_decorator
+    async def similarity_search_by_embedding(
+        self,
+        embedding: list[float],
+        agent_id: int,
+        k: int = 4,
+        fetch_k: int = 20,
+        return_score_type: Union[Literal["none"], Literal["L2-distance"]] = "none",
+        filter: Optional[dict] = None,
+    ) -> Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]:
+        """
+        Return content most similar to the given query.
+
+        Args:
+            embedding (list[float]): The vector to look up documents similar to.
+            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
+            k (int, optional): The number of top similar contents to return. Defaults to 4.
+            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            return_score_type (Union[Literal["none"], Literal["similarity_score"], Literal["L2-distance"]], optional):
+                Specifies whether and how to return similarity scores with the results:
+                    - "none": Do not return scores; only return the contents (default).
+                    - "L2-distance": Return a tuple of content and its L2 distance from the query.
+            filter (dict, optional): The filter dict for metadata.
+
+        Returns:
+            Union[list[tuple[str,dict]], list[tuple[str, float,dict]]]:
+                Depending on the `return_score_type` parameter, returns either a list of strings representing the top-k similar contents,
+                or a list of tuples where each tuple contains a string and a floating-point score.
+        """
+        _filter = {
+            "_id": agent_id,
+        }
+        if filter is not None:
+            _filter.update(filter)
+        if return_score_type == "L2-distance":
+            _result = await self.vectors_store.asimilarity_search_with_score_by_vector(
+                embedding=embedding,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+            )
+            return [(r.page_content, s, r.metadata) for r, s in _result]
+        elif return_score_type == "none":
+            _result = await self.vectors_store.asimilarity_search_by_vector(
+                embedding=embedding,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+            )
+            return [(r.page_content, r.metadata) for r in _result]
+        else:
+            raise ValueError(f"Invalid `return_score_type` {return_score_type}!")
+
+    @lock_decorator
+    async def marginal_relevance_search(
+        self,
+        query: str,
+        agent_id: int,
+        k: int = 4,
+        fetch_k: int = 20,
+        lambda_mult: float = 0.5,
+        return_score_type: Literal["none"] = "none",
+        filter: Optional[dict] = None,
+    ) -> list[tuple[str, dict]]:
+        """
+        Return contents selected using the maximal marginal relevance asynchronously.
+
+        Args:
+            query (str): The text to look up documents similar to.
+            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
+            k (int, optional): The number of top similar contents to return. Defaults to 4.
+            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
+            return_score_type (Literal["none"].,optional):
+                Specifies whether and how to return similarity scores with the results:
+                    - "none": Do not return scores; only return the contents (default).
+            filter (dict, optional): The filter dict for metadata.
+
+        Returns:
+            list[tuple[str,dict]]: the result contents.
+        """
+        _filter = {
+            "_id": agent_id,
+        }
+        if filter is not None:
+            _filter.update(filter)
+
+        if return_score_type == "none":
+            _result = await self.vectors_store.amax_marginal_relevance_search(
+                query=query,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+                lambda_mult=lambda_mult,
+            )
+            return [(r.page_content, r.metadata) for r in _result]
+        else:
+            raise ValueError(f"Invalid `return_score_type` {return_score_type}!")
+
+    @lock_decorator
+    async def marginal_relevance_search_by_embedding(
+        self,
+        embedding: list[float],
+        agent_id: int,
+        k: int = 4,
+        fetch_k: int = 20,
+        lambda_mult: float = 0.5,
+        return_score_type: Union[Literal["none"], Literal["similarity_score"]] = "none",
+        filter: Optional[dict] = None,
+    ) -> Union[list[tuple[str, dict]], list[tuple[str, float, dict]]]:
+        """
+        Return contents selected using the maximal marginal relevance asynchronously.
+
+        Args:
+            embedding (list[float]): The vector to look up documents similar to.
+            agent_id (int): The identifier of the agent to filter specific documents. Only documents associated with this agent will be considered.
+            k (int, optional): The number of top similar contents to return. Defaults to 4.
+            fetch_k (int, optional): The number of documents to fetch before applying any filters. Defaults to 20.
+            lambda_mult (float): Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5.
+            return_score_type (Union[Literal["none"], Literal["similarity_score"]], optional):
+                Specifies whether and how to return similarity scores with the results:
+                    - "none": Do not return scores; only return the contents (default).
+                    - "similarity_score": Return a tuple of content and its similarity score.
+            filter (dict, optional): The filter dict for metadata.
+
+        Returns:
+            Union[list[tuple[str,dict]], list[tuple[str, float,dict]]]:
+                Depending on the `return_score_type` parameter, returns either a list of strings representing the top-k similar contents,
+                or a list of tuples where each tuple contains a string and a floating-point score.
+        """
+
+        _filter = {
+            "_id": agent_id,
+        }
+        if filter is not None:
+            _filter.update(filter)
+        if return_score_type == "none":
+            _result = await self.vectors_store.amax_marginal_relevance_search_by_vector(
+                embedding=embedding,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+                lambda_mult=lambda_mult,
+            )
+            return [(r.page_content, r.metadata) for r in _result]
+        elif return_score_type == "similarity_score":
+            _result = await self.vectors_store.amax_marginal_relevance_search_with_score_by_vector(
+                embedding=embedding,
+                k=k,
+                filter=_filter,
+                fetch_k=fetch_k,
+                lambda_mult=lambda_mult,
+            )
+            return [(r.page_content, s, r.metadata) for r, s in _result]
+
+        else:
+            raise ValueError(f"Invalid `return_score_type` {return_score_type}!")