| from typing import Dict, Any, Optional |
| import os |
| import time |
| import json |
| from openai import OpenAI |
|
|
| class BaseAgent: |
| """Base class for all agents with common functionality.""" |
| |
| def __init__(self, model): |
| """Initialize the base agent with the specified model.""" |
| print(f"Initializing Base Agent with {model}") |
| if model is None: |
| model = os.getenv("DEFAULT_LLM_MODEL") |
| if model is None: |
| raise ValueError("No model specified and DEFAULT_LLM_MODEL not set") |
| |
| openai_key = os.getenv("OPENAI_API_KEY") |
| if openai_key is None: |
| raise ValueError("OPENAI_API_KEY environment variable not set") |
| |
| self.model = model |
| print("Init Client") |
| self.client = OpenAI(api_key="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") |
| print("Client successfully initialized") |
| |
| def __call__(self, question: str, file_name: str = None, **kwargs) -> str: |
| """Process the question and return an answer.""" |
| raise NotImplementedError("Subclasses must implement this method") |
|
|
| def extract_search_query(self, question: str) -> str: |
| """ |
| Extract the key search terms from the question. |
| |
| Args: |
| question: The original question |
| |
| Returns: |
| A reformulated search query |
| """ |
| system_prompt = """ |
| You are an expert at extracting search terms from questions. |
| Given a question, identify the key terms that would be most effective for a web search. |
| Remove any unnecessary words and focus on specific names, dates, or key concepts. |
| Return only the search terms, with no additional text. |
| """ |
| |
| user_prompt = f"Extract the most effective search terms from this question: {question}" |
| |
| response = self.llm_call(system_prompt, user_prompt) |
| return response.strip() |
| |
| def llm_call(self, system_prompt: str, user_prompt: str, |
| temperature: float = 0.1, model: Optional[str] = None) -> str: |
| """ |
| Make a call to the LLM with the given prompts. |
| |
| Args: |
| system_prompt: The system prompt to guide the LLM behavior |
| user_prompt: The user prompt containing the question or task |
| temperature: Controls randomness (lower is more deterministic) |
| model: Optional override for the model to use |
| |
| Returns: |
| The LLM response content as a string |
| """ |
| if model is None: |
| model = self.model |
| |
| try: |
| response = self.client.chat.completions.create( |
| model=model, |
| messages=[ |
| {"role": "system", "content": system_prompt}, |
| {"role": "user", "content": user_prompt} |
| ], |
| temperature=temperature |
| ) |
| return response.choices[0].message.content |
| except Exception as e: |
| print(f"Error calling LLM: {str(e)}") |
| return f"Error: {str(e)}" |
| |
| def ask_james(self, question) -> str: |
| |
| system_prompt = """ |
| You are a general AI Assistant. |
| You must provide only the exact answer requested by the user, in the correct type and minimal format, based purely on the provided information. |
| |
| YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. |
| If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. |
| If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. |
| If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. |
| If you cannot find the answer based on the provided information, reply only with "Unknown". |
| |
| You must not explain your answer, restate the question, or add any commentary. |
| |
| Output must be a single line, without punctuation or extra whitespace. |
| """ |
| final_answer = self.llm_call(system_prompt, question, temperature=0.0) |
|
|
| return final_answer |
