给智能体
外挂两种脑子

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 示例文档
documents = [
    "The sky is blue and beautiful.",
    "Love this blue and beautiful sky!",
    "The quick brown fox jumps over the lazy dog.",
    "A king's breakfast has sausages, ham, bacon, eggs, toast, and beans",
    "I love green eggs, ham, sausages and bacon!",
    "The brown fox is quick and the blue dog is lazy!",
    "The sky is very blue and the sky is very beautiful today",
    "The dog is lazy but the brown fox is quick!"
]

# 第一步：用 TF-IDF 把文档变成向量
vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(documents)

# 第二步：把向量存进一个极简的「向量数据库」（这里就是个列表）
vector_database = X.toarray()

# 第三步：余弦相似度检索
def cosine_similarity_search(query, database, vectorizer, top_n=5):
    query_vec = vectorizer.transform([query]).toarray()
    similarities = cosine_similarity(query_vec, database)[0]
    top_indices = np.argsort(-similarities)[:top_n]  # 取前 n 个下标
    return [(idx, similarities[idx]) for idx in top_indices]

# 交互式检索循环
while True:
    query = input("输入检索词（输入 exit 退出）：")
    if query.lower() == 'exit':
        break
    top_n = int(input("想看前几条结果？"))
    search_results = cosine_similarity_search(query, vector_database, vectorizer, top_n)

    print("最匹配的文档：")
    for idx, score in search_results:
        print(f"- {documents[idx]} (相似度: {score:.4f})")

    print("\n")

import uuid
from pathlib import Path

import chromadb
import tiktoken
from agents import (Agent, Runner, function_tool,
                    set_default_openai_api, set_tracing_disabled)
from dotenv import load_dotenv

# 从 .env 读取 OPENAI_BASE_URL、OPENAI_API_KEY、OPENAI_DEFAULT_MODEL
load_dotenv()
# 下面两行仓库原文没有，是为了接第三方模型加的
set_default_openai_api("chat_completions")
set_tracing_disabled(True)

# ------------------------------------------------------------------
# 1. 读入剧本并切块
#    注意：这是相对路径，必须从仓库根目录运行
# ------------------------------------------------------------------
script_text = Path("chapter_06/sample_documents/back_to_the_future.txt").read_text(
    encoding="utf-8"
)


def simple_chunk(text, max_tokens=200):
    tokenizer = tiktoken.get_encoding("cl100k_base")
    words, chunk, chunks = text.split(), [], []
    for w in words:
        if len(tokenizer.encode(" ".join(chunk + [w]))) > max_tokens:
            chunks.append(" ".join(chunk))
            chunk = [w]
        else:
            chunk.append(w)
    if chunk:
        chunks.append(" ".join(chunk))
    return chunks


docs = simple_chunk(script_text, max_tokens=200)

# ------------------------------------------------------------------
# 2. 建一个 Chroma 集合（或连上已有的）
#    仓库原注释说这里用的是 OpenAI embeddings —— 那是错的，
#    没传 embedding function，用的是 Chroma 自带的本地模型
# ------------------------------------------------------------------
client = chromadb.PersistentClient(
    path="./chapter_06/chroma_script_store"  # 存在磁盘上，之后可以复用
)
collection_name = "bttf_script"

# 先试着拿已有的集合
try:
    collection = client.get_collection(collection_name)
except Exception:
    # 集合不存在就建一个（不指定 embedding function）
    collection = client.create_collection(name=collection_name)

# 只灌一次数据（已经有内容就跳过）
if collection.count() == 0:
    collection.add(ids=[str(uuid.uuid4()) for _ in docs], documents=docs)


# ------------------------------------------------------------------
# 3. 定义一个查 Chroma 的语义检索工具
# ------------------------------------------------------------------
@function_tool
def search_script(query: str, top_k: int = 3) -> str:
    res = collection.query(query_texts=[query], n_results=top_k)
    if res and "documents" in res and res["documents"] and res["documents"][0]:
        return "\n\n".join(res["documents"][0])  # 把最相关的几块拼起来
    return "No relevant documents found."


# ------------------------------------------------------------------
# 4. 组装智能体
# ------------------------------------------------------------------
agent = Agent(
    name="Script Agent",
    instructions=(
        "你负责回答关于电影《回到未来》的问题。\n"
        "需要的时候，调用 `search_script` 工具取回原文片段，"
        "然后在回答里引用或转述它们。"
    ),
    tools=[search_script],
)

# ------------------------------------------------------------------
# 5. 提问
#    查询保持英文 —— 语料是英文剧本，中文查询会检索不准，见下方说明
# ------------------------------------------------------------------
query = "Where does Doc tell Marty to meet him, and at what time?"
result = Runner.run_sync(agent, query)
print("\n--- 回答 ---\n", result.final_output)

query = "What happens at 1:15AM"
result = Runner.run_sync(agent, query)
print("\n--- 回答 ---\n", result.final_output)

async def main():
    # 图记忆：知识图谱
    memory_srv = MCPServerStdio(
        name="memory",
        params={
            "command": "npx",
            "args": ["-y", "@modelcontextprotocol/server-memory@latest"],
        },
    )

    # 语义记忆：ChromaDB 向量库
    # 注意这里用的是 uvx，不是 npx —— chroma-mcp 是个 Python 包
    chroma_srv = MCPServerStdio(
        name="chroma",
        params={
            "command": "uvx",
            "args": ["chroma-mcp", "--data-dir", "chapter_06/chroma_script_store"],
        },
    )

    instructions = """
你是一个混合记忆管理智能体，把语义记忆（ChromaDB 向量库）
和图记忆（知识图谱）结合起来，提供完整的记忆能力。

关键：每一次对话你都必须先说「Remembering...」，
然后执行混合记忆检索流程。

混合记忆系统：
1. 语义记忆（ChromaDB）：存对话内容、文档和上下文信息
2. 图记忆（知识图谱）：存结构化的事实、关系、实体和观察
    """
    # …… 后面还有很长一段提示词，规定两种记忆各自什么时候用 ……

    async with memory_srv, chroma_srv:
        agent = Agent(
            name="Hybrid Memory Agent",
            instructions=instructions,
            mcp_servers=[memory_srv, chroma_srv],
        )
        # …… 交互循环，同上 ……