ehsanulhaque92/multimodal-prescriptive-pdm

1import os2import logging3from langchain_community.vectorstores import Chroma4from langchain_community.embeddings import HuggingFaceEmbeddings5from langchain_community.llms import Ollama6from langchain.prompts import PromptTemplate7from langchain.schema.runnable import RunnablePassthrough8from langchain.schema.output_parser import StrOutputParser9import app_config as config10 11logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')12 13def create_rag_chain():14    embedding_model = HuggingFaceEmbeddings(model_name="./embedding_model")15    vector_store = Chroma(persist_directory=str(config.DB_PATH), embedding_function=embedding_model)16    retriever = vector_store.as_retriever(search_kwargs={"k": 3})17    18    # --- NEW: Always use Ollama, but the URL and model change for deployment ---19    deployment_platform = os.getenv("DEPLOYMENT_PLATFORM", "local")20    21    if deployment_platform == "huggingface":22        # When deployed, Ollama runs in the same container.23        llm_model_name = "gemma:2b"24        base_url = "http://127.0.0.1:11434"25        logging.info(f"Using containerized Ollama with model: {llm_model_name}")26    else:27        # For local Docker, connect to the host machine's Ollama28        llm_model_name = os.getenv("LLM_MODEL_NAME", "llama3:8b")29        base_url = "http://host.docker.internal:11434"30        logging.info(f"Using host Ollama with model: {llm_model_name}")31    32    try:33        llm = Ollama(model=llm_model_name, base_url=base_url)34    except Exception as e:35        logging.error(f"Failed to initialize Ollama: {e}", exc_info=True)36        return None37 38    prompt_template_str = """39    Answer the question based only on the following context.40    Context: {context}41    Question: {question}42    Answer:43    """44    prompt = PromptTemplate.from_template(prompt_template_str)45    rag_chain = (46        {"context": retriever, "question": RunnablePassthrough()}47        | prompt48        | llm49        | StrOutputParser()50    )51    logging.info("RAG chain created successfully.")52    return rag_chain53 54 55if __name__ == '__main__':56    # This block remains for local testing57    print("--- Testing the RAG Chain (Local Ollama Mode) ---")58    try:59        chain = create_rag_chain()60        if chain:61            test_question = "How do I replace the bearing assembly?"62            print(f"\n[?] Test Question: {test_question}")63            answer = chain.invoke(test_question)64            print("\n[!] AI-Generated Answer:")65            print(answer)66            print("\n--- RAG Chain test complete. ---")67        else:68            print("Chain initialization failed.")69    except Exception as e:70        logging.error(f"An error occurred during the RAG chain test: {e}", exc_info=True)