Update app.py

app.py

@@ -1,25 +1,123 @@
-from fastapi import FastAPI
+from langchain_core.prompts import PromptTemplate
+from langchain.chains import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+# import gradio as gr
+# from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import numpy as np
+from langchain_ollama import OllamaLLM
+from langchain_huggingface import HuggingFaceEmbeddings
+# from langchain_community.llms import HuggingFacePipeline
+# from load_document import load_data
+# from split_document import split_docs
+# from embed_docs import embed_docs
+# from retrieve import retrieve
+from datetime import datetime
+# from js import js
+# from theme import theme
+import os
+import glob
+from fastapi import FastAPI, Query, Request
 from pydantic import BaseModel
-from langchain_community.llms import Ollama
-from langchain.callbacks.manager import CallbackManager
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+import uvicorn
 
-app = FastAPI()
-MODEL_NAME = 'tinyllama'
 
-def get_llm():
-    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
-    return Ollama(model=MODEL_NAME, callback_manager=callback_manager)
+app = FastAPI(title="Know The Law", description="A FastAPI application for legal assistance using AI.")
 
-class Question(BaseModel):
-    text: str
+
+
+vector_store_path = "/home/user/VectorStoreDB"
+index_name = "faiss_index"
+full_index_path = os.path.join(vector_store_path, index_name)
+
+# # Create the embedder with a specific model
+embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+# # Initialize our speech pipeline
+# transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base.en", device="cpu")
+
+
+
+
+def fetch_doc():
+    # Adjust the path as needed, e.g., './' for current directory
+    pdf_files = glob.glob("Document/*.pdf")
+    return pdf_files
+
+# # Define llm
+hf_token = os.environ.get("HF_TOKEN").strip()  # Ensure to set your Hugging Face token in the environment variable HF_TOKEN
+llm = OllamaLLM(model="mistral:7b-instruct")
+
+
+pdf_files = fetch_doc() #Fetch Dataset
+chunks = None
+loaded_docs = []
+# just query if it exists
+if not os.path.exists(full_index_path):
+    for doc in pdf_files:
+        print(f"Loading.....{doc}")
+        docs = load_data(doc) #Load Dataset
+        loaded_docs.append(docs)
+    final_docs = [item for sublist in loaded_docs for item in sublist] # Flatten the list    
+    chunks = split_docs(final_docs, embedder=embedder) #Split Document
+saved_vector = embed_docs(chunks, embedder=embedder) #Embed Document
+retrieved = retrieve(saved_vector) # Retrieve simimlar docs
+
+# Define the prompt template
+prompt = """
+You are The Law Assistant, an AI trained to help Nigerians understand their legal rights and obligations. Using the provided context below, answer user questions related to Nigerian law.
+
+Instructions:
+
+1. Base your responses strictly on the given context or verified legal sources.
+
+2. If the answer is not in the context and you're unsure, respond with: "I don't know based on the available information." Do not fabricate or speculate.
+
+3. Keep your answers clear, concise, and jargon-free.
+
+4. Always cite the legal source(s) or reference(s) you used (e.g., constitution section, legal act, court ruling).
+
+Context: {context}
+Question: {{question}}
+
+Helpful Answer:"""
+
+
+QA_CHAIN_PROMPT = PromptTemplate.from_template(template=prompt)
+
+# Create document prompt
+document_prompt = PromptTemplate(
+    input_variables=["page_content", "source"],
+    template="Context:\ncontent:{page_content}\nsource:{source}",
+)
+
+# Create the stuff documents chain
+combine_docs_chain = create_stuff_documents_chain(
+    llm,
+    QA_CHAIN_PROMPT,
+    document_prompt=document_prompt
+)
+
+# Create the retrieval chain
+qa_chain = create_retrieval_chain(
+    retriever=retrieved,
+    combine_docs_chain=combine_docs_chain
+)
+
+class QueryRequest(BaseModel):
+    question: str
 
 @app.get("/")
-def read_root():
-    return {"message": f"Welcome to {MODEL_NAME} FastAPI"}
-
-@app.post("/ask")
-def ask_question(question: Question):
-    llm = get_llm()
-    response = llm(question.text)
-    return {"response": response}
+def home():
+    return {"message": "Welcome to the Know The Law API. Use POST /query to ask legal questions."}
+
+@app.post("/query")
+def respond(query: QueryRequest):
+    # Invoke the chain with the question
+    question = query.question
+    result = qa_chain.invoke({"input":question})
+    
+    # Return the answer
+    return {"answer": result['answer']}
+
+
+