from langchain_community.llms.ollama import Ollama
# from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate


PROMPT_TEMPLATE = """Examples

Document: USC is one of the top universities in America
Topic: Universities
Answer: Yes

Document: USC is one of the top universities in America
Topic: Politics
Answer: No

Document: Governments worldwide are grappling with the challenges of regulating artificial intelligence to ensure it benefits society while protecting individual privacy.
Topic: Politics
Answer: Yes

Answer if the Topic Classification of the following Document is correct. 
Answer only in "Yes" or "No".

Document: {document} 
Topic: {topic}
Answer:"""



def topic_in_text_service(document: str, topic: str, model_id: str, base_url: str) -> bool:
    llm = Ollama(model="llama3", base_url="http://ollama:11434")
    prompt_template = PromptTemplate.from_template(PROMPT_TEMPLATE)

    prompt = prompt_template.format(document=document, topic=topic)

    return 'yes' in llm.invoke(prompt, num_predict=1).lower()