From d9d826fae5671af0137dcdb41a9b1ac61ece1761 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=205f1b81b8-4f5d-11e8-9c2d-fa7ae01?= =?UTF-8?q?bbebc?= Date: Thu, 2 Jul 2026 16:45:49 +0000 Subject: [PATCH] Updated main.py with Ollama embeddings and LangChain agent --- main.py | 214 ++++++++++++++++++-------------------------------------- 1 file changed, 70 insertions(+), 144 deletions(-) diff --git a/main.py b/main.py index 3594815..f062844 100644 --- a/main.py +++ b/main.py @@ -1,162 +1,88 @@ import asyncio -import os -import json from pathlib import Path -from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_chroma import Chroma -from langchain_core.documents import Document -from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_core.messages import HumanMessage +from langchain_ollama import ChatOllama, OllamaEmbeddings +from langchain.embeddings import OpenAIEmbeddings +from langchain.vectorstores import Chroma +from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType +from langchain.tools import BaseTool +import httpx +import os -# ------------------------------------------------------------------ -# Configuration -# ------------------------------------------------------------------ -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") -if not OPENAI_API_KEY: - raise RuntimeError("OPENAI_API_KEY not set in environment") +# Load FAQ data +DATA_DIR = Path("data") -# ------------------------------------------------------------------ -# LLM and embeddings (OpenRouter only) -# ------------------------------------------------------------------ -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, - temperature=0.0, -) +# Embedding model +embeddings = OllamaEmbeddings(model="nomic-embed-text") -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, -) +# Create Chroma store +def load_faq_to_chroma() -> Chroma: + from langchain.document_loaders import TextLoader + from langchain.text_splitter import RecursiveCharacterTextSplitter -# ------------------------------------------------------------------ -# Chroma vector store (persisted) -# ------------------------------------------------------------------ -CHROMA_PATH = Path("./chroma_faq") -vector_store = Chroma( - collection_name="faq_collection", - embedding_function=embeddings, - persist_directory=str(CHROMA_PATH), -) - -# ------------------------------------------------------------------ -# Utility: load markdown files into Chroma -# ------------------------------------------------------------------ -async def load_faq_to_chroma(md_dir: str = "data"): - md_dir = Path(md_dir) - if not md_dir.is_dir(): - raise FileNotFoundError(f"Markdown directory {md_dir} not found") docs = [] - for md_file in md_dir.glob("*.md"): - text = md_file.read_text(encoding="utf-8") - docs.append(Document(page_content=text, metadata={"source": md_file.name})) - vector_store.add_documents(docs) - vector_store.persist() - print(f"Loaded {len(docs)} documents into Chroma (persisted at {CHROMA_PATH})") + for md_file in DATA_DIR.glob("*.md"): + loader = TextLoader(str(md_file)) + docs.extend(loader.load_and_split(RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50))) + db = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_faq") + db.persist() + return db -# ------------------------------------------------------------------ -# Tools -# ------------------------------------------------------------------ -@tool -def search_course_docs(query: str) -> str: - """Search the local FAQ collection for relevant passages.""" - results = vector_store.similarity_search(query, k=3) - if not results: - return "No relevant information found in the course materials." - return "\n---\n".join(f"[{doc.metadata.get('source', 'unknown')}] {doc.page_content}" for doc in results) +chroma_db = load_faq_to_chroma() -@tool -def fetch_course_meta(query: str) -> str: - """Simulate an MCP-style tool that fetches course metadata. - In production this would be a real HTTP call to an MCP server. - Here we use a local JSON file as a mock response. - """ - meta_file = Path("meta.json") - if not meta_file.is_file(): - return "Metadata source not available." - data = json.loads(meta_file.read_text(encoding="utf-8")) - # Very naive search: return items where query string appears in any value - matches = [] - for key, value in data.items(): - if isinstance(value, str) and query.lower() in value.lower(): - matches.append(f"{key}: {value}") - if not matches: - return "No metadata matches found." - return "\n".join(matches) +# Tool: search in FAQ +class SearchFAQTool(BaseTool): + name = "search_course_docs" + description = "Search local FAQ docs. Use query string. Returns top k results." -# ------------------------------------------------------------------ -# Backend setup -# ------------------------------------------------------------------ -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) + def _run(self, query: str, k: int = 3): + results = chroma_db.similarity_search_with_score(query, k) + return "\n".join([f"{i+1}. {r[0].page_content[:200]}... (score: {r[1]:.4f})" for i, r in enumerate(results)]) -# ------------------------------------------------------------------ -# DeepAgent creation -# ------------------------------------------------------------------ -agent = create_deep_agent( - model=llm, - tools=[search_course_docs, fetch_course_meta], - backend=backend, - system_prompt=( - "You are a helpful FAQ assistant for the course. " - "When answering a question, use the local Chroma database if the answer is about course content. " - "If the question is about schedule, metadata, or other non-content info, call fetch_course_meta. " - "Always indicate the source in your final answer as either 'source: chroma' or 'source: mcp_meta'." - ), +search_tool = SearchFAQTool() + +# Tool: fetch course metadata (MCP style) +class FetchMetaTool(BaseTool): + name = "fetch_course_meta" + description = "Fetch course metadata via HTTP. Use query string. Returns JSON string." + + def _run(self, query: str): + # For demo, use local JSON file or mock endpoint + url = f"http://localhost:8000/meta?query={query}" + try: + resp = httpx.get(url, timeout=5) + resp.raise_for_status() + return resp.text + except Exception as e: + return f"Error fetching meta: {e}" + +meta_tool = FetchMetaTool() + +# LLM +llm = ChatOllama(model="llama3") + +# Agent +tools = [search_tool, meta_tool] + +agent = initialize_agent( + tools, + llm, + agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, + verbose=True, + handle_parsing_errors=True, ) -# ------------------------------------------------------------------ -# CLI helpers -# ------------------------------------------------------------------ -PRESET_QUESTIONS = [ - "What topics are covered in the first lecture?", - "Explain the concept of recursion as described in the notes.", - "When is the next lab session scheduled?", -] - -async def run_interactive(): - print("--- FAQ Bot CLI ---") - print("Type 'exit' to quit.") - while True: - user_input = input("\nQuestion: ") - if user_input.lower() in {"exit", "quit"}: - break - result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - # The last message is the assistant's reply - reply = result["messages"][-1].content - print("\nAnswer:\n", reply) - -async def run_presets(): - for q in PRESET_QUESTIONS: - print("\nQuestion:", q) - result = await agent.ainvoke( - {"messages": [HumanMessage(content=q)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - reply = result["messages"][-1].content - print("Answer:\n", reply) - -# ------------------------------------------------------------------ -# Main entry point -# ------------------------------------------------------------------ async def main(): - # Load data into Chroma if not already persisted - if not CHROMA_PATH.is_dir() or not any(CHROMA_PATH.iterdir()): - await load_faq_to_chroma() - # Run preset questions first - await run_presets() - # Then interactive mode - await run_interactive() + # Simple CLI with predefined questions + questions = [ + "What is the deadline for assignment 3?", + "How to use ChromaDB with LangChain?", + "What is the schedule for next week?", + ] + for q in questions: + print("\nQuestion:", q) + result = await agent.arun(input=q) + print("Answer:\n", result) if __name__ == "__main__": asyncio.run(main())