diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..4208e80 --- /dev/null +++ b/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()