|
|
|
|
|
|
|
|
import os |
|
|
import logging |
|
|
from typing import List, Dict, Any, Optional |
|
|
from abc import ABC, abstractmethod |
|
|
from langchain_openai import ChatOpenAI |
|
|
from langchain.schema import SystemMessage, HumanMessage, AIMessage |
|
|
|
|
|
from ..config.config import ( |
|
|
OPENAI_API_KEY, |
|
|
AGENT_DEFAULT_TEMPERATURE, |
|
|
AGENT_DEFAULT_MODEL, |
|
|
MAX_TOKENS |
|
|
) |
|
|
|
|
|
class BaseAgent(ABC): |
|
|
"""Base class for all agents in the AI Co-Scientist system.""" |
|
|
|
|
|
def __init__(self, |
|
|
name: str, |
|
|
system_prompt: str, |
|
|
model: str = AGENT_DEFAULT_MODEL, |
|
|
temperature: float = AGENT_DEFAULT_TEMPERATURE): |
|
|
"""Initialize the base agent. |
|
|
|
|
|
Args: |
|
|
name: The name of the agent |
|
|
system_prompt: The system prompt that defines the agent's role |
|
|
model: The LLM model to use |
|
|
temperature: The temperature parameter for generation |
|
|
""" |
|
|
self.name = name |
|
|
self.system_prompt = system_prompt |
|
|
|
|
|
|
|
|
self.model = model if model is not None else AGENT_DEFAULT_MODEL |
|
|
self.temperature = temperature |
|
|
self.conversation_history = [] |
|
|
|
|
|
|
|
|
self.llm = ChatOpenAI( |
|
|
model=self.model, |
|
|
temperature=self.temperature, |
|
|
api_key=OPENAI_API_KEY, |
|
|
max_tokens=MAX_TOKENS |
|
|
) |
|
|
|
|
|
|
|
|
self.conversation_history.append( |
|
|
SystemMessage(content=system_prompt) |
|
|
) |
|
|
|
|
|
self.logger = logging.getLogger(f"agent.{name}") |
|
|
|
|
|
def add_message(self, message: str, is_human: bool = True) -> None: |
|
|
"""Add a message to the conversation history. |
|
|
|
|
|
Args: |
|
|
message: The message content |
|
|
is_human: Whether the message is from a human (True) or AI (False) |
|
|
""" |
|
|
if is_human: |
|
|
self.conversation_history.append(HumanMessage(content=message)) |
|
|
else: |
|
|
self.conversation_history.append(AIMessage(content=message)) |
|
|
|
|
|
def get_response(self, query: str) -> str: |
|
|
"""Get a response from the agent based on the query. |
|
|
|
|
|
Args: |
|
|
query: The query to send to the agent |
|
|
|
|
|
Returns: |
|
|
The agent's response as a string |
|
|
""" |
|
|
|
|
|
self.add_message(query, is_human=True) |
|
|
|
|
|
|
|
|
response = self.llm.invoke(self.conversation_history) |
|
|
|
|
|
|
|
|
self.add_message(response.content, is_human=False) |
|
|
|
|
|
return response.content |
|
|
|
|
|
def clear_history(self) -> None: |
|
|
"""Clear the conversation history, keeping only the system prompt.""" |
|
|
system_prompt = self.conversation_history[0] |
|
|
self.conversation_history = [system_prompt] |
|
|
|
|
|
@abstractmethod |
|
|
def process(self, input_data: Any) -> Any: |
|
|
"""Process the input data and return a result. |
|
|
|
|
|
This method must be implemented by all subclasses. |
|
|
|
|
|
Args: |
|
|
input_data: The input data to process |
|
|
|
|
|
Returns: |
|
|
The processing result |
|
|
""" |
|
|
pass |
|
|
|
