chat.py

import requests

from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings

DB_PATH = "vector_db"

embedding = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

db = Chroma(
    persist_directory=DB_PATH,
    embedding_function=embedding
)

def ask_llm(context, question):
    prompt = f"""
You MUST answer ONLY using the provided context.

If the answer is not in the context, say:
"Not found in provided documents."

Context:
{context}

Question:
{question}

Answer grounded strictly in the context:
"""

    response = requests.post(
        "http://localhost:11434/api/generate",
        json={
            "model": "mistral",
            "prompt": prompt,
            "stream": False
        }
    )

    return response.json()["response"]

while True:
    q = input("You: ")
    docs = db.similarity_search(q, k=6)
    context = "\n".join([d.page_content for d in docs])
    print("AI:", ask_llm(context, q))