Model.py

from dotenv import load_dotenv
from augmented import load_db, embedding_model
import os
from openai import OpenAI
load_dotenv()
class OpenRouterLLM:
    def __init__(self):
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=os.getenv("OPENROUTER_API_KEY"),
        )
        self.model = os.getenv("LLM_MODEL")

        if not self.model:
            raise ValueError("LLM_MODEL not set in .env")

    def invoke(self, prompt: str):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
        )

        # mimic LangChain response object
        class Response:
            def __init__(self, content):
                self.content = content

        return Response(response.choices[0].message.content)
llm=OpenRouterLLM()

def model(question):
    os.makedirs("data", exist_ok=True)
    if len(os.listdir("data/")) == 0:
        return "No Data Was Found and Please upload the data"
    else:
        db = load_db()
        query_embed = embedding_model.embed_query(question)
        result=db.similarity_search_by_vector(embedding=query_embed,k=3)
        context = "\n\n".join([doc.page_content for doc in result])
        prompt = f"""
            Answer the question using only the context below:
            {context}

            Question: {question}
            """
        response=llm.invoke(prompt)
    return response