Please fix my code:

import os
import logging
import time
import psutil
import asyncio
import uuid
import pickle
import torch
import numpy as np
from typing import Optional, Dict
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from langchain_core.prompts import PromptTemplate
from langchain_community.llms import Ollama
from transformers import AutoTokenizer
from sentence_transformers import SentenceTransformer
import faiss

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# Initialize the FastAPI app
app = FastAPI()

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize SentenceTransformer model and tokenizer
model_name = 'all-MiniLM-L6-v2'
sentence_model = SentenceTransformer(model_name)
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# Define the persist directory and normalize the path
persist_dir = os.path.normpath("D:/Storage")

# Load document embeddings, metadata, and FAISS index
with open(os.path.join(persist_dir, 'document_embeddings.pkl'), 'rb') as f:
    document_embeddings = pickle.load(f)

with open(os.path.join(persist_dir, 'document_metadata.pkl'), 'rb') as f:
    document_metadata = pickle.load(f)

index = faiss.read_index(os.path.join(persist_dir, 'faiss_index.bin'))

# Create a mapping of document IDs to their metadata
document_metadata_dict = {doc['id']: doc['extra_info'] for doc in document_metadata}

# Define a custom RAG prompt
custom_rag_prompt = """
You are a helpful assistant. Provide a clear, structured, and concise response to the following question using the given context.
Ensure your answer is quick and efficient, and correct. Include detailed steps if needed. Do not repeat statements in your answer.

Question: {question}

Answer:
{context}

Note: The information was retrieved from the following sources:
{source_info}
"""

# Create a PromptTemplate with the custom RAG prompt
custom_rag_prompt_template = PromptTemplate(
    template=custom_rag_prompt,
    input_variables=["question", "context", "source_info"]
)

# Function to format documents and extract source information
def format_docs_and_sources(docs):
    formatted_docs = []
    source_info = []

    for doc in docs:
        if hasattr(doc, 'text') and hasattr(doc, 'metadata'):
            content = f"Table:\n{doc.text}" if "table" in doc.id_ else doc.text
            formatted_docs.append(content)

            doc_metadata = getattr(doc, 'metadata', {})
            doc_name = doc_metadata.get('name', 'Unknown Document Name')
            pages_info = doc_metadata.get('pages', 'No page information available')
            source_info.append(f"{doc_name}, Pages: {pages_info}")
        else:
            formatted_docs.append("No text available")
            source_info.append("No source information available")

    formatted_docs = "\n\n".join(formatted_docs)
    source_info = "\n".join(list(set(source_info)))

    return formatted_docs, source_info

# Function to log server performance
def log_performance():
    cpu_usage = psutil.cpu_percent(interval=1)
    memory_info = psutil.virtual_memory()
    disk_io = psutil.disk_io_counters()

    logging.info(f"CPU Usage: {cpu_usage}%")
    logging.info(f"Memory Usage: {memory_info.percent}% (Available: {memory_info.available / (1024 * 1024):.2f}MB)")
    logging.info(f"Disk Read: {disk_io.read_bytes / (1024 * 1024):.2f}MB, Disk Write: {disk_io.write_bytes / (1024 * 1024):.2f}MB")

# Dictionary to manage background tasks and their statuses
tasks: Dict[str, dict] = {}
websockets: Dict[str, WebSocket] = {}

# Define the request body model
class QARequest(BaseModel):
    question: str
    model: str = "stablelm-zephyr"
    temperature: float = 0.7
    top_p: float = 0.1
    top_k: int = 50
    client: str  # Added client field

# Function to encode text using SentenceTransformer
def encode_text(text: str) -> np.ndarray:
    return sentence_model.encode(text, convert_to_numpy=True)

async def process_qa(question: str, model: str = "stablelm-zephyr", temperature: float = 0.7, top_p: float = 0.1, top_k: int = 50, task_id: Optional[str] = None, client: str = ""):
    try:
        start_time = time.time()  # Start time for the entire process
        logging.info(f"Process QA started for task_id: {task_id}")
        logging.info(f"Received request with question: '{question}' and model: '{model}'")

        # Encode the question
        logging.info("Encoding the question")
        question_embedding = encode_text(question)

        # Normalize client name for comparison
        client_lower = client.lower().strip()
        logging.info(f"Normalized client name: {client_lower}")

        # Create a mapping of client names to their indices
        client_name_mapping = {doc['extra_info']['client'].lower().strip(): idx for idx, doc in enumerate(document_metadata)}

        # Get the indices for the client
        client_filtered_indices = [idx for c, idx in client_name_mapping.items() if c == client_lower]

        if not client_filtered_indices:
            logging.info(f"No documents found for client: {client}")
            tasks[task_id] = {"status": "Completed", "result": "No relevant information found for the specified client."}
            await send_result_to_websocket(task_id, "No relevant information found for the specified client.", 0)
            return "No relevant information found for the specified client.", "", 0

        # Create a filtered FAISS index for the specified client
        filtered_index = faiss.IndexFlatL2(sentence_model.get_sentence_embedding_dimension())  # Use the new model's dimension
        filtered_embeddings = np.array([document_embeddings[idx] for idx in client_filtered_indices], dtype='float32')
        filtered_index.add(filtered_embeddings)

        # Perform FAISS search on the filtered index
        logging.info("Performing FAISS search for shorter summaries")
        faiss_result = filtered_index.search(np.array([question_embedding], dtype='float32'), k=3)  # Get top 3 documents

        if isinstance(faiss_result, tuple) and len(faiss_result) == 2:
            distances, indices = faiss_result
        else:
            logging.error(f"Unexpected result from FAISS search: {faiss_result}")
            raise ValueError(f"Unexpected result from FAISS search: {faiss_result}")

        logging.info(f"Top 3 document indices: {indices[0]} with distances: {distances[0]}")

        # Retrieve relevant shorter summary documents
        shorter_summary_docs = [document_metadata[client_filtered_indices[idx]] for idx in indices[0]]
        
        # Now get the longer summaries from the top shorter summaries
        longer_summary_docs = []
        for doc in shorter_summary_docs:
            # Assume you have a method to get the longer summary
            longer_summary = doc.get("longer_summary", doc.text)  # Use the longer summary if available
            longer_summary_docs.append(longer_summary)

        # Encode longer summaries and get the best one
        longer_summary_embeddings = np.array([encode_text(doc) for doc in longer_summary_docs])
        longer_summary_index = faiss.IndexFlatL2(sentence_model.get_sentence_embedding_dimension())
        longer_summary_index.add(longer_summary_embeddings)
        
        faiss_result_longer = longer_summary_index.search(np.array([question_embedding], dtype='float32'), k=1)  # Get top document
        best_longer_summary_index = faiss_result_longer[1][0][0]
        best_longer_summary_doc = longer_summary_docs[best_longer_summary_index]

        # Prepare the context for the LLM
        formatted_context, source_info = format_docs_and_sources([best_longer_summary_doc])

        # Prepare the prompt for the language model
        inputs = {
            "context": formatted_context,
            "question": question,
            "source_info": source_info
        }
        prompt = custom_rag_prompt_template.format(**inputs)

        # Call the language model with the prompt
        try:
            model_start_time = time.time()
            llm = Ollama(model=model, temperature=temperature, top_p=top_p, top_k=top_k)
            model_output = await asyncio.get_event_loop().run_in_executor(None, llm.call, prompt)
            model_time = time.time() - model_start_time
            logging.info(f"Model inference completed in {model_time:.2f} seconds")
        except Exception as e:
            logging.error(f"Error during model inference: {e}")
            raise

        # Log the response from the model
        logging.info(f"Model output: {model_output}")

        # Update the task status and return the result
        tasks[task_id] = {"status": "Completed", "result": model_output}
        await send_result_to_websocket(task_id, model_output, model_time)

        return model_output, source_info, model_time

    except Exception as e:
        logging.error(f"Error occurred in process_qa: {e}")
        if task_id:
            tasks[task_id] = {"status": "Error", "result": str(e)}
        raise

# Function to handle WebSocket connections
async def send_result_to_websocket(task_id: str, result: str, model_time: float):
    if task_id in websockets:
        websocket = websockets[task_id]
        try:
            await websocket.send_json({
                "task_id": task_id,
                "status": "completed",
                "result": result,
                "model_time": model_time
            })
        except WebSocketDisconnect:
            logging.warning(f"WebSocket disconnected for task_id: {task_id}")

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    task_id = str(uuid.uuid4())
    websockets[task_id] = websocket
    tasks[task_id] = {"status": "Running", "result": None}

    try:
        while True:
            data = await websocket.receive_json()
            qa_request = QARequest(**data)

            # Start processing the request in the background
            asyncio.create_task(
                process_qa(
                    question=qa_request.question,
                    model=qa_request.model,
                    temperature=qa_request.temperature,
                    top_p=qa_request.top_p,
                    top_k=qa_request.top_k,
                    task_id=task_id,
                    client=qa_request.client
                )
            )

            await websocket.send_json({
                "task_id": task_id,
                "status": "submitted"
            })

    except WebSocketDisconnect:
        logging.warning(f"WebSocket disconnected for task_id: {task_id}")
        del websockets[task_id]

# Root endpoint for testing
@app.get("/")
async def root():
    return {"message": "FastAPI RAG system is running!"}


INFO:     Started server process [6220]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
2024-09-26 08:44:26,909 - WARNING - WebSocket disconnected for task_id: fd592ad1-9922-43ae-86e1-395f9830d49f
INFO:     connection closed
INFO:     ('172.21.18.79', 53120) - "WebSocket /ws" [accepted]
INFO:     connection open
2024-09-26 08:44:35,176 - INFO - Process QA started for task_id: 39ed7aba-c4e3-4ebf-b66f-aeb72307760d
2024-09-26 08:44:35,176 - INFO - Received request with question: 'When does the contract end' and model: 'qwen2:0.5b'
2024-09-26 08:44:35,176 - INFO - Encoding the question
Batches: 100%|███████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 21.13it/s]
2024-09-26 08:44:35,223 - INFO - Normalized client name: absa
2024-09-26 08:44:35,223 - ERROR - Error occurred in process_qa:
2024-09-26 08:44:35,223 - ERROR - Task exception was never retrieved
future: <Task finished name='Task-9' coro=<process_qa() done, defined at C:\Users\FourieH\API_RAG_cont_A-DLR.py:126> exception=AssertionError()>
Traceback (most recent call last):
  File "C:\Users\FourieH\API_RAG_cont_A-DLR.py", line 155, in process_qa
    filtered_index.add(filtered_embeddings)
  File "C:\Users\FourieH\AppData\Local\anaconda3\Lib\site-packages\faiss\class_wrappers.py", line 228, in replacement_add
    assert d == self.d
           ^^^^^^^^^^^
AssertionError