Updated main.py with Ollama embeddings and LangChain agent

main.py

@@ -1,162 +1,88 @@
 import asyncio
-import os
-import json
 from pathlib import Path
 
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_ollama import ChatOllama, OllamaEmbeddings
-from langchain_chroma import Chroma
+from langchain.embeddings import OpenAIEmbeddings
-from langchain_core.documents import Document
+from langchain.vectorstores import Chroma
-from langchain.tools import tool
+from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType
-from deepagents import create_deep_agent
+from langchain.tools import BaseTool
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+import httpx
-from langchain_core.messages import HumanMessage
+import os
 
-# ------------------------------------------------------------------
+# Load FAQ data
-# Configuration
+DATA_DIR = Path("data")
-# ------------------------------------------------------------------
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-if not OPENAI_API_KEY:
-    raise RuntimeError("OPENAI_API_KEY not set in environment")
 
-# ------------------------------------------------------------------
+# Embedding model
-# LLM and embeddings (OpenRouter only)
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
-# ------------------------------------------------------------------
-llm = ChatOpenAI(
-    model="openai/gpt-oss-20b:free",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=OPENAI_API_KEY,
-    temperature=0.0,
-)
 
-embeddings = OpenAIEmbeddings(
+# Create Chroma store
-    model="text-embedding-3-small",
+def load_faq_to_chroma() -> Chroma:
-    base_url="https://openrouter.ai/api/v1",
+    from langchain.document_loaders import TextLoader
-    api_key=OPENAI_API_KEY,
+    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"):
+    for md_file in DATA_DIR.glob("*.md"):
-        text = md_file.read_text(encoding="utf-8")
+        loader = TextLoader(str(md_file))
-        docs.append(Document(page_content=text, metadata={"source": md_file.name}))
+        docs.extend(loader.load_and_split(RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)))
-    vector_store.add_documents(docs)
+    db = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_faq")
-    vector_store.persist()
+    db.persist()
-    print(f"Loaded {len(docs)} documents into Chroma (persisted at {CHROMA_PATH})")
+    return db
 
-# ------------------------------------------------------------------
+chroma_db = load_faq_to_chroma()
-# 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)
 
-@tool
+# Tool: search in FAQ
-def fetch_course_meta(query: str) -> str:
+class SearchFAQTool(BaseTool):
-    """Simulate an MCP-style tool that fetches course metadata.
+    name = "search_course_docs"
-    In production this would be a real HTTP call to an MCP server.
+    description = "Search local FAQ docs. Use query string. Returns top k results."
-    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)
 
-# ------------------------------------------------------------------
+    def _run(self, query: str, k: int = 3):
-# Backend setup
+        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)])
-backend = CompositeBackend([
-    LocalShellBackend(workspace_dir="./workspace"),
-    FilesystemBackend(),
-])
 
-# ------------------------------------------------------------------
+search_tool = SearchFAQTool()
-# DeepAgent creation
+
-# ------------------------------------------------------------------
+# Tool: fetch course metadata (MCP style)
-agent = create_deep_agent(
+class FetchMetaTool(BaseTool):
-    model=llm,
+    name = "fetch_course_meta"
-    tools=[search_course_docs, fetch_course_meta],
+    description = "Fetch course metadata via HTTP. Use query string. Returns JSON string."
-    backend=backend,
+
-    system_prompt=(
+    def _run(self, query: str):
-        "You are a helpful FAQ assistant for the course. "
+        # For demo, use local JSON file or mock endpoint
-        "When answering a question, use the local Chroma database if the answer is about course content. "
+        url = f"http://localhost:8000/meta?query={query}"
-        "If the question is about schedule, metadata, or other non-content info, call fetch_course_meta. "
+        try:
-        "Always indicate the source in your final answer as either 'source: chroma' or 'source: mcp_meta'."
+            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
+    # Simple CLI with predefined questions
-    if not CHROMA_PATH.is_dir() or not any(CHROMA_PATH.iterdir()):
+    questions = [
-        await load_faq_to_chroma()
+        "What is the deadline for assignment 3?",
-    # Run preset questions first
+        "How to use ChromaDB with LangChain?",
-    await run_presets()
+        "What is the schedule for next week?",
-    # Then interactive mode
+    ]
-    await run_interactive()
+    for q in questions:
+        print("\nQuestion:", q)
+        result = await agent.arun(input=q)
+        print("Answer:\n", result)
 
 if __name__ == "__main__":
     asyncio.run(main())