from typing import TypedDict, Annotated, List
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, AIMessage

class AgentState(TypedDict):
    # 对话历史（工作记忆的核心）
    messages: Annotated[List, add_messages]
    # 当前任务的中间结果
    task_context: dict
    # 已执行的步骤记录
    steps_taken: List[str]
    # 最终结果
    final_result: str

def create_agent_graph():
    graph = StateGraph(AgentState)

    def think_node(state: AgentState) -> AgentState:
        # 工作记忆在整个执行过程中持续传递
        messages = state["messages"]
        context = state.get("task_context", {})
        steps = state.get("steps_taken", [])

        # 模型能看到完整的对话历史和任务状态
        print(f"当前已执行步骤：{steps}")
        print(f"工作记忆中的消息数：{len(messages)}")

        return state

    graph.add_node("think", think_node)
    graph.add_edge(START, "think")
    graph.add_edge("think", END)

    return graph.compile()

from openai import OpenAI
import chromadb
import json
from datetime import datetime

client = OpenAI()
chroma_client = chromadb.Client()
collection = chroma_client.get_or_create_collection(name="agent_episodes")

def save_episode(task: str, approach: str, result: str, quality_score: float):
    """保存一次任务执行经历到情节记忆"""
    episode = {
        "task": task,
        "approach": approach,
        "result_summary": result,
        "quality_score": quality_score,
        "timestamp": datetime.now().isoformat()
    }

    # 用任务描述作为向量化内容，方便语义检索
    collection.add(
        documents=[task],
        metadatas=[episode],
        ids=[f"episode_{datetime.now().timestamp()}"]
    )
    print(f"已保存执行经历，质量评分：{quality_score}")

def retrieve_relevant_episodes(current_task: str, n_results: int = 3) -> List[dict]:
    """检索与当前任务相似的历史经历"""
    results = collection.query(
        query_texts=[current_task],
        n_results=n_results
    )

    episodes = []
    if results["metadatas"]:
        for metadata in results["metadatas"][0]:
            episodes.append(metadata)

    return episodes

def agent_with_episodic_memory(task: str) -> str:
    # 先检索相关历史经验
    past_episodes = retrieve_relevant_episodes(task)

    memory_context = ""
    if past_episodes:
        memory_context = "以下是处理类似任务的历史经验：\n"
        for ep in past_episodes:
            score = ep.get("quality_score", 0)
            approach = ep.get("approach", "")
            # 只参考质量高的经验
            if score >= 0.7:
                memory_context += f"- 任务类型：{ep['task'][:50]}...\n"
                memory_context += f"  采用方案：{approach}\n"
                memory_context += f"  质量评分：{score}\n\n"

    messages = [
        {
            "role": "system",
            "content": f"你是一个有经验的分析助手。{memory_context}"
        },
        {"role": "user", "content": task}
    ]

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages
    )
    result = response.choices[0].message.content

    # 执行完成后，保存这次经历
    save_episode(
        task=task,
        approach="直接分析 + 历史经验参考",
        result=result[:200],
        quality_score=0.8  # 实际场景应该由 Critic 评分
    )

    return result

import chromadb
from openai import OpenAI

client = OpenAI()
chroma_client = chromadb.Client()
knowledge_base = chroma_client.get_or_create_collection(name="semantic_memory")

def add_knowledge(content: str, source: str, category: str):
    """向知识库添加内容"""
    knowledge_base.add(
        documents=[content],
        metadatas=[{"source": source, "category": category}],
        ids=[f"doc_{hash(content)}"]
    )

def retrieve_knowledge(query: str, n_results: int = 3) -> str:
    """从知识库检索相关内容"""
    results = knowledge_base.query(
        query_texts=[query],
        n_results=n_results
    )

    if not results["documents"] or not results["documents"][0]:
        return ""

    retrieved_docs = []
    for doc, metadata in zip(results["documents"][0], results["metadatas"][0]):
        retrieved_docs.append(f"[来源：{metadata['source']}]\n{doc}")

    return "\n\n---\n\n".join(retrieved_docs)

# 预先加载公司内部知识
add_knowledge(
    content="代码审查时，关注以下高风险模式：1. SQL 拼接（SQL 注入风险）2. 未加密的敏感数据存储 3. 未验证的用户输入直接传递给命令行",
    source="安全规范文档 v2.1",
    category="security"
)

def rag_agent(question: str) -> str:
    # 从语义记忆检索相关知识
    relevant_knowledge = retrieve_knowledge(question)

    messages = [
        {
            "role": "system",
            "content": f"你是一个代码审查专家。\n\n参考知识：\n{relevant_knowledge}"
        },
        {"role": "user", "content": question}
    ]

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages
    )
    return response.choices[0].message.content

def smart_memory_retrieval(query: str, min_quality: float = 0.7) -> list:
    """结合语义检索和精确过滤"""
    # 第一步：语义检索，找相似内容
    results = collection.query(
        query_texts=[query],
        n_results=10  # 先多拿一些
    )

    if not results["metadatas"]:
        return []

    # 第二步：精确过滤，只保留高质量的
    filtered = [
        metadata for metadata in results["metadatas"][0]
        if metadata.get("quality_score", 0) >= min_quality
    ]

    # 按质量评分排序
    filtered.sort(key=lambda x: x.get("quality_score", 0), reverse=True)

    return filtered[:3]  # 最多返回 3 条