import subprocess

import os

from pathlib import Path

from langgraph.graph import StateGraph, END

from langchain_ollama import ChatOllama, OllamaEmbeddings

import chromadb

from typing import TypedDict, List, Optional

import requests

import json

import re

# from langchain.text_splitter import RecursiveCharacterTextSplitter

from assets.split import split_text

from tools.weather.weather_tool import WeatherTool

# ===================== LLM =====================

llm = ChatOllama(model="llama3.2:1b", temperature=0.1)

ollama_embeddings = OllamaEmbeddings(

model="nomic-embed-text",

base_url="http://127.0.0.1:11434"

)

# ===================== 记忆库 + RAG向量库 =====================

chroma_client = chromadb.PersistentClient(path="./agent_memory")

# 原有记忆库

memory_collection = chroma_client.get_or_create_collection(name="agent_long_term_memory")

# 新增RAG文件库

rag_collection = chroma_client.get_or_create_collection(name="rag_document_store")

# ===================== Ollama Embedding =====================

def get_ollama_embedding(text: str) -> list[float]:

try:

embedding = ollama_embeddings.embed_query(text)

return embedding

except Exception as e:

raise Exception(f"Embedding 失败：{e}")

# ===================== 状态定义（新增RAG相关字段） =====================

class AgentState(TypedDict):

user_input: str

city: str

execution_result: str

retrieved_memory: Optional[str]

decision: Optional[str]

# 新增RAG字段

file_path: Optional[str] # 上传的文件路径

retrieved_rag_content: Optional[str] # 检索到的文件内容

# ===================== 天气工具 =====================

weather_tool = WeatherTool()

# ===================== 原有记忆逻辑 =====================

MEMORY_KEYWORDS = ["记住", "保存", "记一下", "记录"]

def need_memory(user_input: str) -> bool:

return any(k in user_input for k in MEMORY_KEYWORDS)

def store_memory(query: str, result: str):

embedding = get_ollama_embedding(query)

memory_collection.upsert(

ids=[f"mem_{hash(query)}"],

embeddings=[embedding],

metadatas=[{"query": query, "result": result}]

)

def retrieve_memory(query: str, threshold=0.75) -> Optional[str]:

embedding = get_ollama_embedding(query)

res = memory_collection.query(query_embeddings=[embedding], n_results=1)

if not res["metadatas"][0]:

return None

sim = 1 - res["distances"][0][0]

return res["metadatas"][0][0]["result"] if sim >= threshold else None

# ===================== 新增RAG核心逻辑 =====================

# 1. 文件解析函数

def parse_file(file_path: str) -> str:

"""解析本地文件，提取纯文本内容"""

try:

file_path = Path(file_path).resolve()

# 安全校验：仅允许访问当前目录

if not file_path.is_relative_to(Path.cwd()):

raise ValueError("禁止访问上级目录")

# 支持的文件类型

ext = file_path.suffix.lower()

if ext in [".txt", ".md", ".py", ".json", ".yml", ".yaml"]:

with open(file_path, "r", encoding="utf-8") as f:

return f.read()

else:

raise ValueError(f"不支持的文件类型：{ext}，仅支持txt/md/py/json/yml/yaml")

except Exception as e:

return f"文件解析失败：{str(e)}"

# 2. 上传文件到RAG库

def upload_file_to_rag(file_path: str) -> str:

"""上传文件并分块存入RAG向量库"""

# 解析文件内容

raw_text = parse_file(file_path)

if "失败" in raw_text:

return raw_text

# 文本分块

chunks = split_text(raw_text)

if not chunks:

return "文件内容为空"

# 生成向量并入库

try:

file_name = Path(file_path).name

for i, chunk in enumerate(chunks):

chunk_id = f"rag_{file_name}_{i}_{hash(chunk)}"

embedding = get_ollama_embedding(chunk)

rag_collection.upsert(

ids=[chunk_id],

embeddings=[embedding],

metadatas=[{"file_name": file_name, "chunk_index": i, "content": chunk}]

)

return f"✅ 文件上传成功！\n文件名：{file_name}\n分块数量：{len(chunks)}"

except Exception as e:

return f"RAG入库失败：{str(e)}"

# 3. RAG检索函数

def retrieve_rag_content(query: str, threshold=0.7) -> str:

"""检索与问题相关的文件内容"""

embedding = get_ollama_embedding(query)

res = rag_collection.query(query_embeddings=[embedding], n_results=3) # 取Top3相关片段

# 过滤低相似度结果

relevant_chunks = []

for meta, dist in zip(res["metadatas"][0], res["distances"][0]):

sim = 1 - dist

if sim >= threshold and meta:

relevant_chunks.append({

"file_name": meta.get("file_name", "未知文件"),

"content": meta.get("content", ""),

"similarity": round(sim, 2)

})

if not relevant_chunks:

return ""

# 拼接检索结果

rag_content = "📄 检索到相关文件内容：\n"

for i, chunk in enumerate(relevant_chunks, 1):

rag_content += f"\n【{i}】文件：{chunk['file_name']}（相似度：{chunk['similarity']}）\n{chunk['content']}\n"

return rag_content

# 4. RAG问答生成

def generate_rag_answer(state: AgentState) -> str:

"""基于检索到的文件内容生成回答"""

query = state["user_input"]

rag_content = state["retrieved_rag_content"]

prompt = f"""

基于以下文件内容回答用户问题，仅使用文件中的信息，不要编造内容：

{rag_content}

用户问题：{query}

回答要求：

1. 简洁明了，基于文件内容回答；

2. 如果文件中没有相关信息，直接说"未在上传的文件中找到相关信息"；

3. 不要提及"文件"、"片段"等词汇，自然回答即可。

"""

try:

res = llm.invoke(prompt)

return res.content.strip()

except Exception as e:

return f"RAG回答生成失败：{str(e)}"

# ===================== 原有工具函数 =====================

def run_command(cmd: str) -> str:

try:

r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=20)

return f"命令结果：\n{r.stdout}" if r.returncode == 0 else f"失败：{r.stderr}"

except Exception as e:

return f"执行出错：{str(e)}"

def read_file(path_str: str) -> str:

try:

p = Path(path_str).resolve()

if not p.is_relative_to(Path.cwd()):

return "禁止访问上级目录"

with open(p, "r", encoding="utf-8") as f:

return f"文件内容：\n{f.read()[:1500]}"

except Exception as e:

return f"读取失败：{e}"

# ===================== 分类函数（新增rag分类） =====================

def classify_input(state: AgentState) -> AgentState:

user_input = state["user_input"]

# 先查看向量数据库里面是否有答案

ok , result = has_answer_in_rag(user_input)

if ok:

state["decision"] = "rag_qa"

return state

# 上传rag

if "上传文件" in user_input:

state["decision"] = "rag_upload"

state["file_path"] = user_input.split(" ")[1].strip(" ")

print(state["file_path"])

return state

prompt = f"""

# 任务要求（必须100%遵守）

你需要仅返回一个JSON字符串，格式固定为：{{"category":"分类值","city":"城市名"}}，无任何其他内容。

# 分类规则

1. category 可选值：weather/command/file/chat/rag_upload/rag_qa（仅这6个，无其他）；

2. weather：所有天气相关请求，city提取用户输入中的城市名（无则为空字符串）；

3. command：执行系统命令的请求，city固定为空字符串；

4. file：读取/查看文件内容的请求（非上传），city固定为空字符串；

5. chat：日常聊天内容，city固定为空字符串；

6. rag_upload：上传文件到知识库的请求，注意读的是本地的文件，没让你上传文件（包含"上传文件 文件路径"、"导入文件 文件路径"、"rag上传 文件路径"等关键词）；

7. rag_qa：基于已上传文件提问的请求（包含"根据文件"、"文件里"、"知识库"等关键词）；

8. 不确定分类时，category默认值为chat，city为空字符串。

# 禁止项

- 禁止返回嵌套JSON、多余字段；

- 禁止返回代码块、解释文字、换行、空格；

- 禁止修改JSON字段名。

# 用户输入

{user_input}

# 输出结果（仅JSON字符串，无其他）：

"""

try:

res = llm.invoke(prompt)

result = res.content.strip().lower() # 统一用content（新版langchain）

if result.startswith("```"):

result = result.replace("```json", "").replace("```", "").strip()

result_json = json.loads(result)

print("result: ", result_json)

# 解析分类结果

state["decision"] = result_json["category"]

state["city"] = result_json["city"]

# 提取上传文件路径（如果是rag_upload）

# if state["decision"] == "rag_upload":

# # 从用户输入中提取文件路径（如：上传文件 ./test.txt → ./test.txt）

# path_pattern = r"([./\\][\w\-.\\/]+)"

# path_match = re.search(path_pattern, user_input)

# state["file_path"] = path_match.group(1) if path_match else ""

# else:

# state["file_path"] = None

return state

except Exception as e:

print("分类失败:", e)

# 异常时默认chat

state["decision"] = "chat"

state["city"] = ""

return state

# ===================== 节点（新增RAG相关节点） =====================

# 原有节点

def memory_node(state: AgentState) -> AgentState:

res = retrieve_memory(state["user_input"])

if res:

state["execution_result"] = f"【记忆】{res}"

return state

def chat_node(state: AgentState) -> AgentState:

user = state["user_input"]

prompt = f"你是智能助手，自然回答：{user}"

try:

res = llm.invoke(prompt)

ans = res.content.strip()

state["execution_result"] = ans

if need_memory(user):

store_memory(user, ans)

state["execution_result"] += "\n\n【已记住这条回答】"

return state

except Exception as e:

state["execution_result"] = f"聊天回答失败：{str(e)}"

return state

def command_node(state: AgentState) -> AgentState:

res = run_command(state["user_input"])

state["execution_result"] = res

return state

def file_node(state: AgentState) -> AgentState:

res = read_file(state["user_input"])

state["execution_result"] = res

return state

def weather_node(state: AgentState) -> AgentState:

res = weather_tool.get_weather(state["city"])

state["execution_result"] = res

return state

# 新增RAG节点

def rag_upload_node(state: AgentState) -> AgentState:

"""上传文件到RAG库"""

file_path = state["file_path"]

if not file_path:

state["execution_result"] = "❌ 未识别到文件路径，请输入格式：上传文件 ./test.txt"

return state

# 执行上传

upload_result = upload_file_to_rag(file_path)

state["execution_result"] = upload_result

return state

def rag_qa_node(state: AgentState) -> AgentState:

"""基于RAG库问答"""

# 1. 检索相关文件内容

rag_content = retrieve_rag_content(state["user_input"])

state["retrieved_rag_content"] = rag_content

# 2. 生成回答

if not rag_content:

state["execution_result"] = "未在上传的文件中找到相关信息"

else:

answer = generate_rag_answer(state)

state["execution_result"] = answer

return state

# ===================== 路由（新增RAG分支） =====================

def router_after_memory(state: AgentState) -> str:

d = state["decision"]

if d == "chat": return "chat"

if d == "command": return "command"

if d == "file": return "file"

if d == "weather": return "weather"

if d == "rag_upload": return "rag_upload" # 新增

if d == "rag_qa": return "rag_qa" # 新增

return "chat"

# 从向量数据库拿到答案

def has_answer_in_rag(question: str, threshold: float = 0.7) -> tuple[bool, str]:

"""

判断知识库中是否存在问题的答案

:param question: 用户问题

:param threshold: 相似度阈值（0-1，默认0.7）

:return: (是否存在答案, 检索到的相关内容/无答案原因)

"""

try:

# 1. 生成问题的向量

query_embedding = get_ollama_embedding(question)

# 2. 检索向量库（取Top3最相似片段）

res = rag_collection.query(

query_embeddings=[query_embedding],

n_results=3,

include=["metadatas", "distances"] # 需获取元数据和距离

)

# 3. 无检索结果

if not res["metadatas"][0] or len(res["distances"][0]) == 0:

return False, "未检索到任何相关文件内容"

# 4. 过滤低相似度结果（相似度=1-距离）

relevant_chunks = []

for meta, dist in zip(res["metadatas"][0], res["distances"][0]):

similarity = 1 - dist

if similarity >= threshold and meta:

relevant_chunks.append({

"file_name": meta.get("file_name", "未知文件"),

"content": meta.get("content", ""),

"similarity": round(similarity, 2)

})

# 5. 无有效相似片段

if not relevant_chunks:

avg_similarity = round(1 - sum(res["distances"][0])/len(res["distances"][0]), 2)

return False, f"检索到的内容相似度过低（平均相似度：{avg_similarity}），低于阈值 {threshold}"

# 6. 存在有效答案，拼接相关内容

content_str = "📄 找到相关答案（来自以下文件）：\n"

for i, chunk in enumerate(relevant_chunks, 1):

content_str += f"\n【{i}】文件：{chunk['file_name']}（相似度：{chunk['similarity']}）\n{chunk['content']}\n"

return True, content_str

except Exception as e:

return False, f"判断失败：{str(e)}"

# ===================== 构建图（新增RAG节点和边） =====================

def build():

g = StateGraph(AgentState)

# 原有节点

g.add_node("classify_input", classify_input)

g.add_node("chat", chat_node)

g.add_node("command", command_node)

g.add_node("file", file_node)

g.add_node("weather", weather_node)

# 新增RAG节点

g.add_node("rag_upload", rag_upload_node)

g.add_node("rag_qa", rag_qa_node)

g.set_entry_point("classify_input")

# 条件边（新增rag_upload/rag_qa）

g.add_conditional_edges("classify_input", router_after_memory, {

"chat": "chat",

"command": "command",

"file": "file",

"weather": "weather",

"rag_upload": "rag_upload",

"rag_qa": "rag_qa",

END: END

})

# 所有节点指向结束

g.add_edge("chat", END)

g.add_edge("command", END)

g.add_edge("file", END)

g.add_edge("weather", END)

g.add_edge("rag_upload", END) # 新增

g.add_edge("rag_qa", END) # 新增

return g.compile()

# ===================== 运行 =====================

if __name__ == "__main__":

agent = build()

print("=== 智能AI助手（聊天+命令+文件+天气+记忆+RAG）===")

print("基础功能：聊天、执行命令、读取文件、查询天气、记忆内容")

print("RAG功能：")

print(" - 上传文件：输入「上传文件 ./test.txt」（支持txt/md/py/json等）")

print(" - 文件问答：输入「根据上传的文件回答xxx」「文件里的xxx是什么」")

print("输入 exit 退出，clear memory 清空记忆，clear rag 清空RAG库\n")

while True:

i = input("你：").strip()

if i.lower() == "exit":

break

# 清空记忆

if i.lower() == "clear memory":

ids = memory_collection.get()["ids"]

if ids: memory_collection.delete(ids=ids)

print("✅ 记忆已清空\n")

continue

# 新增：清空RAG库

if i.lower() == "clear rag":

ids = rag_collection.get()["ids"]

if ids: rag_collection.delete(ids=ids)

print("✅ RAG文件库已清空\n")

continue

# 调用Agent

res = agent.invoke({

"user_input": i,

"execution_result": "",

"retrieved_memory": None,

"decision": None,

"file_path": None,

"retrieved_rag_content": None

})

print("AI：" + res["execution_result"] + "\n")