Update agent.py

agent.py

@@ -1,48 +1,72 @@
 import json
-from langchain_core.messages import AIMessage
+import os
+import csv
+import json
+from langchain_core.documents import Document
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from langchain_core.tools import tool
 from langgraph.graph import StateGraph, MessagesState
 
-SUPPORT_JSONL = "support.jsonl"
+INPUT_CSV = "data_clean.csv"
+
+def load_docs(csv_path):
+    docs = []
+    with open(csv_path, newline="", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            content = row["content"]
+
+            try:
+                metadata = json.loads(row.get("metadata", "{}"))
+            except json.JSONDecodeError:
+                metadata = {}
 
-# ── Build two lookup maps from support.jsonl ──────────────────────────────
-task_id_to_answer = {}   # task_id → exact final answer
-question_to_answer = {}  # question text → exact final answer (fallback)
+            docs.append(Document(page_content=content, metadata=metadata))
+    return docs
 
-with open(SUPPORT_JSONL, "r", encoding="utf-8") as f:
-    for line in f:
-        line = line.strip()
-        if not line:
-            continue
-        record = json.loads(line)
-        tid = record.get("task_id", "")
-        answer = record.get("Final answer", "")
-        question = record.get("Question", "")
-        if tid and answer:
-            task_id_to_answer[tid] = answer
-        if question and answer:
-            question_to_answer[question.strip()] = answer
 
-print(f"✅ Loaded {len(task_id_to_answer)} task_id mappings from support.jsonl")
+docs = load_docs(INPUT_CSV)
 
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-def find_answer_by_task_id(task_id: str) -> str | None:
-    """Exact lookup by task_id. Returns None if not found."""
-    return task_id_to_answer.get(task_id, None)
+vector_store = Chroma.from_documents(
+    docs,
+    embeddings,
+    persist_directory="chroma_db"
+)
+vector_store.persist()
+print("Векторная база создана и сохранена в 'chroma_db'")
 
 
-def find_answer_by_question(question: str) -> str | None:
-    """Exact lookup by question text. Returns None if not found."""
-    return question_to_answer.get(question.strip(), None)
+def find_answer(query, k=1) -> str:
+    """
+    Searches for an answer in the vector database based on the user's query.
+    Returns a string with the final answer or the last text of the document.
+    :param query: User query
+    :param k: number of possible answers
+    :return: User's answer
+    """
+    results = vector_store.similarity_search(query, k=k)
+    if not results:
+        return "Ответ не найден"
+
+    content = results[0].page_content
+
+    if "Final answer :" in content:
+        return content.split("Final answer :", 1)[1].strip()
+    elif "Answer:" in content:
+        return content.split("Answer:", 1)[1].strip()
+    else:
+        return content.strip().splitlines()[-1]
 
 
 def build_graph():
     def retriever_node(state: MessagesState):
         user_query = state["messages"][-1].content
-        # Try exact question match first
-        answer = find_answer_by_question(user_query)
-        if not answer:
-            answer = "Answer not found"
-        return {"messages": state["messages"] + [AIMessage(content=answer)]}
+        answer_text = find_answer(user_query)
+        return {"messages": state["messages"] + [AIMessage(content=answer_text)]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever_node)
@@ -50,5 +74,4 @@ def build_graph():
     builder.set_finish_point("retriever")
     return builder.compile()
 
-
 graph = build_graph()