from langchain_core.language_models.llms import LLM
from typing import Any, List, Optional, Union
import subprocess
import requests
import time
import json
import re

PromptType = Union[str, List[str], dict]

#1.Ollama
class Ollama(LLM):
  @property
  def _llm_type(self)-> str:
    return "Ollama open source models without installations"

  def _call(self, prompts: PromptType, model: str, stop: Optional[List[str]] = None) -> str:
    text_prompt = self._format_prompt(prompts)
    try:
      url = "https://npmai-api.onrender.com/llm"
      payload = {"prompt": text_prompt,"model": model,}
      response = requests.post(url, json=payload)
      return response.json()["response"]
     
    except Exception as e:
      return f"Error during NPMAI Call:str{str(e)}]"

  def invoke(self, prompts: PromptType, model: str) -> str:
    return self._call(prompts, model)

  def _format_prompt(self, prompts: PromptType) -> str:
    if isinstance(prompts, str):
      return prompts
    elif isinstance(prompts, list):
      return "\n".join(str(p) for p in prompts)
    elif isinstance(prompts, dict):
      return json.dumps(prompts)
    else:
      raise ValueError("Invalid prompt type")


# Call Code
if __name__ == "__main__":
    models = {"1": Ollama,}
    choice = input("Choose model 1:-Ollama,").strip()
    prompts = input("Enter query: ").strip()
    Selected = models.get(choice, Ollama)
    llm = Selected()
    response = llm.invoke(prompts)
    print(response)