Add agent.py

agent.py

@@ -0,0 +1,161 @@
+import os
+import json
+import httpx
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+from langchain_ollama import ChatOllama, OllamaEmbeddings
+from langchain_text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.schema import Document
+from langchain import LLMChain, PromptTemplate
+
+# Load environment variables
+load_dotenv()
+
+# Constants
+CHROMA_PATH = Path("chromadb")
+DATA_PATH = Path("data")
+
+# ---------- Data Loader ----------
+
+def load_md_files(data_dir: Path) -> list[Document]:
+    docs = []
+    for file_path in data_dir.rglob("*.md"):
+        try:
+            text = file_path.read_text(encoding="utf-8")
+        except Exception as e:
+            print(f"Error reading {file_path}: {e}")
+            continue
+        docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
+    return docs
+
+
+def chunk_documents(docs: list[Document], chunk_size: int = 1000, chunk_overlap: int = 200) -> list[Document]:
+    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+    all_chunks = []
+    for doc in docs:
+        chunks = splitter.split_text(doc.page_content)
+        for i, chunk in enumerate(chunks):
+            meta = dict(doc.metadata)
+            meta["chunk_index"] = i
+            all_chunks.append(Document(page_content=chunk, metadata=meta))
+    return all_chunks
+
+
+def init_vectorstore() -> Chroma:
+    # If vectorstore exists, load it; otherwise build and persist.
+    if CHROMA_PATH.exists() and any(CHROMA_PATH.iterdir()):
+        return Chroma(persist_directory=str(CHROMA_PATH), embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
+    # Build
+    md_docs = load_md_files(DATA_PATH)
+    if not md_docs:
+        raise RuntimeError("No .md files found in data/")
+    chunks = chunk_documents(md_docs)
+    vectorstore = Chroma.from_documents(chunks, OllamaEmbeddings(model="nomic-embed-text"), persist_directory=str(CHROMA_PATH))
+    return vectorstore
+
+# ---------- MCP Tool ----------
+
+def fetch_course_meta(query: str) -> str:
+    """MCP-style HTTP GET returning JSON.
+    Returns JSON string or error message.
+    """
+    url = "https://api.example.com/course_meta"
+    params = {"query": query}
+    try:
+        response = httpx.get(url, params=params, timeout=10.0)
+        response.raise_for_status()
+        data = response.json()
+        return json.dumps(data, ensure_ascii=False)
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+# ---------- Agent Logic ----------
+SYSTEM_PROMPT = """
+You are a FAQ assistant. Use the provided tools to answer questions.
+- If the question is about course metadata, use the 'fetch_course_meta' tool.
+- For other questions, search the local FAQ database.
+Return the answer and a source field which is either 'chroma' or 'mcp_meta'.
+"""
+
+# LLM
+llm = ChatOllama(model="llama3", base_url=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"))
+
+# Retrieval chain for chroma
+vectorstore = init_vectorstore()
+retriever = vectorstore.as_retriever(search_kwargs={"k":5})
+
+# Prompt template for retrieval
+RETRIEVE_PROMPT = PromptTemplate(
+    input_variables=["question"],
+    template="Answer the question based on the following context. If no context is relevant, answer directly. Context:\n{context}\nQuestion: {question}",
+)
+
+# Retrieval chain (LLM + retriever)
+retrieval_chain = LLMChain(llm=llm, prompt=RETRIEVE_PROMPT)
+
+# Function to get answer from chroma
+def answer_from_chroma(question: str) -> tuple[str, str]:
+    docs = retriever.get_relevant_documents(question)
+    if not docs:
+        # Fallback to MCP
+        meta_json = fetch_course_meta(question)
+        return (meta_json, "mcp_meta")
+    context = "\n".join([doc.page_content for doc in docs])
+    try:
+        answer = retrieval_chain.run({"question": question, "context": context})
+    except Exception as e:
+        answer = f"Error generating answer: {e}"
+    return (answer, "chroma")
+
+# Main query function
+def query_agent(question: str) -> dict:
+    # Simple rule: if question contains 'meta' or 'course', use MCP
+    if any(word in question.lower() for word in ["meta", "course", "information", "details"]):
+        meta_json = fetch_course_meta(question)
+        return {"answer": meta_json, "source": "mcp_meta"}
+    # Else try chroma
+    answer, src = answer_from_chroma(question)
+    return {"answer": answer, "source": src}
+
+# ---------- CLI ----------
+PRESET_QUESTIONS = [
+    "What is the grading policy?",
+    "How do I submit assignments?",
+    "Tell me about course meta for CS101",
+]
+
+def main():
+    print("=== FAQ Bot ===")
+    print("Preset questions:")
+    for i, q in enumerate(PRESET_QUESTIONS, 1):
+        print(f"{i}. {q}")
+    print("Enter 0 to exit or type a custom question.")
+    while True:
+        try:
+            inp = input("\nYour choice (number or question): ")
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting.")
+            break
+        if inp.strip() == "0":
+            print("Goodbye!")
+            break
+        if inp.strip().isdigit():
+            idx = int(inp) - 1
+            if 0 <= idx < len(PRESET_QUESTIONS):
+                question = PRESET_QUESTIONS[idx]
+            else:
+                print("Invalid number.")
+                continue
+        else:
+            question = inp.strip()
+        if not question:
+            print("Please enter a question.")
+            continue
+        result = query_agent(question)
+        print(f"\nAnswer:\n{result['answer']}\nSource: {result['source']}")
+
+if __name__ == "__main__":
+    main()