main.py

import requests
from langchain_core.prompts import ChatPromptTemplate
from vector import retriever
from langchain.llms.base import LLM
from pydantic import Field
import json

# 👉 Replace this with your actual server URL
OLLAMA_URL = "http://localhost:11434/"


class RemoteOllamaLLM(LLM):
    model: str = Field(default="llama3.2")
    url: str = Field(default=OLLAMA_URL)

    def _call(self, prompt, **kwargs):
        response = requests.post(
            f"{self.url}/api/generate",
            json={
                "model": self.model,
                "prompt": prompt,
                "stream": False
            }
        )

        try:
            data = response.json()
        except Exception:
            # Fallback for multiline JSON lines
            first_line = response.text.strip().split("\n")[0]
            data = json.loads(first_line)

        # Check for response key instead of message.content
        if "response" not in data:
            raise Exception(f"Unexpected LLM response format: {response.text}")

        return data["response"]

    @property
    def _llm_type(self) -> str:
        return "remote-ollama"


# Use the remote model
model = RemoteOllamaLLM(model="llama3.2")

# Prompt template (kept for structure, but not used by the model directly)
template = """
You are an expert in answering questions about a pizza restaurant

Here are some relevant reviews: {reviews}

Here is the question to answer: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
chain = prompt | model

# Chat loop
while True:
    print("\n\n-------------------------------")
    question = input("Ask your question (q to quit): ")
    print("\n\n")
    if question.lower() == "q":
        break

    reviews = retriever.invoke(question)
    result = chain.invoke({"reviews": reviews, "question": question})
    print(result)