Allow model and max tokens to be dynamic parameters in Standalone LLM Tool.

This commit is contained in:
cyclop-bot
2024-10-28 10:19:22 -05:00
parent 055399c514
commit 73b6321c5c
+7 -9
View File
@@ -6,26 +6,24 @@ from openai import OpenAI
class StandaloneLLMTool: class StandaloneLLMTool:
def __init__(self): def __init__(self):
self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY")) self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
self.model = "llm-preview"
self.max_tokens = 16384
def get_detailed_instructions(self, user_prompt): def get_detailed_instructions(self, user_prompt, model="llm-preview", max_tokens=16384):
response = self.client.completions.create( response = self.client.completions.create(
model=self.model, model=model,
prompt=user_prompt, prompt=user_prompt,
max_tokens=self.max_tokens max_tokens=max_tokens
) )
return response return response
def process_user_input(self, user_prompt): def process_user_input(self, user_prompt, model="llm-preview", max_tokens=16384):
logging.info(f"Received prompt: {user_prompt}") logging.info(f"Received prompt: {user_prompt}")
response = self.get_detailed_instructions(user_prompt) response = self.get_detailed_instructions(user_prompt, model, max_tokens)
logging.info("Response generated") logging.info("Response generated")
return response.choices[0].text return response.choices[0].text
# Utility function for programmatic access # Utility function for programmatic access
def get_llm_response(prompt): def get_llm_response(prompt, model="llm-preview", max_tokens=16384):
tool = StandaloneLLMTool() tool = StandaloneLLMTool()
return tool.process_user_input(prompt) return tool.process_user_input(prompt, model, max_tokens)