task_6a1d75c5fd30e81cf3126ae7_auto/agent.py

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()