diff --git a/base_agent/llminterface.py b/base_agent/llminterface.py index 7abbf28..fc313f0 100644 --- a/base_agent/llminterface.py +++ b/base_agent/llminterface.py @@ -1,3 +1,4 @@ +import instructor from openai import OpenAI from ollama import Client import ollama @@ -112,3 +113,57 @@ def get_ollama_response(self, question: str, context: str) -> str: self.chat_history.enqueue(response['message']) return response['message']['content'] + +class StructuredLangModel: + """ + Interface to interact with language models using structured query models + """ + + def __init__(self, model: str = 'gpt-4o'): + if 'gpt' in model: + api_key = os.getenv('OPENAI_API_KEY') + self.llm = instructor.from_openai(OpenAI(api_key=api_key)) + else: + self.llm = instructor.from_openai( + OpenAI( + base_url=os.getenv('OLLAMA_HOST', 'http://localhost:11434'), + api_key=os.getenv('OLLAMA_API_KEY', 'ollama') + )) + + def get_response(self, question: str, context: str = None, response_model: object =None) -> str: + if 'gpt' in self.model: + return self.get_gpt_response(question, context,response_model) + else: + return self.get_ollama_response(question, context, response_model) + + def get_gpt_response(self, question: str, context: str, response_model: object) -> str: + self.chat_history.enqueue({'role': 'user', 'content': context + '\n\n'+ question}) + messages=[{'role': 'system', 'content': context}] + self.chat_history.get_all() + + response = self.llm.chat.completions.create(model=self.model, + messages=messages, + response_model=response_model, + max_tokens=1000, + temperature=0.1, top_p=1) + self.chat_history.enqueue(response['message']) + return response.model_dump_json(indent=2) + + def get_ollama_response(self, question: str, context: str, response_model: object) -> str: + """ + Get response from any Ollama supported model + :param question: question to ask + :param context: context to provide + :return: model's response + """ + msg = {'role': 'user', 'content': context + '\n\n' + question} + self.chat_history.enqueue(msg) + messages = self.chat_history.get_all() + response = self.llm.chat.completions.create( + model=self.model, + messages=messages, + response_model= response_model, + options={'temperature': 0} + ) + self.chat_history.enqueue(response['message']) + + return response.model_dump_json(indent=2) \ No newline at end of file