Create main.py

main.py

@@ -0,0 +1,151 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from langchain_ollama import OllamaLLM
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+
+import traceback
+
+# from langchain_core.output_parsers import StrOutputParser
+# from langchain_core.runnables import RunnablePassthrough
+
+import os
+
+os.environ["HF_HOME"] = "/tmp/huggingface"
+
+app = FastAPI()
+
+os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
+
+# Load and split documents
+loader = TextLoader("knowledge_base.txt", encoding="utf-8")
+documents = loader.load()
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=500,
+    chunk_overlap=50,
+    separators=["\n\n", "\n", ".", "!", "?", "،", "؟", "!", ";", ","],
+)
+splits = text_splitter.split_documents(documents)
+
+# Generate embeddings and store in FAISS
+embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
+vectorstore = FAISS.from_documents(splits, embeddings)
+retriever = vectorstore.as_retriever(search_kwargs={"k": 5, "score_threshold": 0.4})
+
+# Define improved prompt template
+template = """
+You are an AI assistant. You must ALWAYS respond in the EXACT SAME LANGUAGE as the user's question or message. This is crucial:
+- If the user writes in English, you MUST respond in English
+- If the user writes in Arabic, you MUST respond in Arabic (Modern Standard Arabic)
+- Mixed language messages should get responses in the predominant language of the message
+Conversation history:
+{history}
+Relevant information from knowledge base:
+{context}
+User's message: {question}
+Key requirements:
+1. MATCH THE LANGUAGE OF THE USER'S MESSAGE EXACTLY
+2. Use the provided context and history to answer the question
+3. Maintain your identity as an AI assistant
+4. Never pretend to be the user or adopt their name
+5. For greetings and casual conversation, respond naturally without using the knowledge base
+6. Only use the knowledge base content when directly relevant to a specific question
+Response:
+"""
+
+prompt = ChatPromptTemplate.from_template(template)
+
+# Load model with adjusted parameters
+model = OllamaLLM(
+    model="mistral",
+    temperature=0.1,
+    num_ctx=8192,
+    top_p=0.8,
+)
+
+
+def format_conversation_history(history):
+    formatted = ""
+    for entry in history:
+        formatted += f"{entry}\n"
+    return formatted
+
+
+# Create RAG chain with properly handled input types
+def generate_response(question, history, retriever):
+    # Get relevant documents
+    context = retriever.invoke(question)
+    context_str = "\n".join(doc.page_content for doc in context)
+
+    # Format the conversation history
+    history_str = format_conversation_history(history)
+
+    # Prepare the input for the prompt
+    chain_input = {"context": context_str, "history": history_str, "question": question}
+
+    # Generate response using the prompt template and model
+    response = prompt.format(**chain_input)
+    response = model.invoke(response)
+
+    return response
+
+
+def chatbot_conversation():
+    print("Hello! I'm an AI assistant. Type 'exit' to quit.")
+
+    conversation_history = []
+
+    while True:
+        user_input = input("You: ").strip()
+        if user_input.lower() == 'exit':
+            break
+
+        try:
+            # Generate response
+            result = generate_response(user_input, conversation_history, retriever)
+
+            print(f"Assistant: {result}")
+
+            # Store the exchange in history
+            conversation_history.append(f"User: {user_input}")
+            conversation_history.append(f"Assistant: {result}")
+
+        except Exception as e:
+            print(f"An error occurred: {str(e)}")
+            print(
+                "Assistant: I apologize, but I encountered an error. Please try again."
+            )
+
+
+chat_histories = {}
+
+
+class ChatRequest(BaseModel):
+    user_id: str  # Unique ID for tracking history per user
+    message: str
+
+
+@app.post("/chat")
+def chat(request: ChatRequest):
+    try:
+        # Retrieve the user's conversation history or create a new one
+        if request.user_id not in chat_histories:
+            chat_histories[request.user_id] = []
+
+        # Get conversation history
+        history = chat_histories[request.user_id]
+
+        # Generate response
+        response = generate_response(request.message, history, retriever)
+
+        # Update history
+        history.append(f"User: {request.message}")
+        history.append(f"Assistant: {response}")
+
+        return {"response": response}
+    except Exception as e:
+        print(traceback.format_exc())
+        raise HTTPException(status_code=500, detail=str(e))