Simulatrex.ai
import asyncio
import dotenv
from simulatrex import SimulationEngine

dotenv.load_dotenv()


async def main():
    engine = SimulationEngine("./data/1_consumer_price_simulation_config.json")
    await engine.run()


if __name__ == "__main__":
    asyncio.run(main())

examples/1_consumer_price_simulation.py
simulatrex/simulatrex
import asyncio
import dotenv
from simulatrex import SimulationEngine

dotenv.load_dotenv()


async def main():
    engine = SimulationEngine("./data/2_policy_impact_local_communities.json")
    await engine.run()


if __name__ == "__main__":
    asyncio.run(main())

examples/2_policy_impact_simulation.py
simulatrex/simulatrex
import asyncio
import dotenv
from simulatrex import SimulationEngine

dotenv.load_dotenv()


async def main():
    engine = SimulationEngine("./data/3_crisis_management_simulation.json")
    await engine.run()


if __name__ == "__main__":
    asyncio.run(main())

examples/3_crisis_management_simulation.py
simulatrex/simulatrex
from setuptools import setup, find_packages

with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

setup(
    name="simulatrex",
    version="0.0.2.1",
    author="Dominik Scherm",
    author_email="me@dominikscherm.de",
    description="LLM-based simulation framework",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/simulatrex/simulatrex",
    packages=find_packages(where="src"),
    include_package_data=True,
    package_dir={"": "src"},
    package_data={
        "simulatrex": ["llm_utils/prompt_templates/*.txt"],
    },
    install_requires=[
        "openai",
        "uuid",
        "numpy",
        "termcolor",
        "chromadb",
        "pydantic",
        "python-dotenv",
        "requests",
        "instructor",
    ],
)

setup.py
simulatrex/simulatrex
from .engine import SimulationEngine

src/simulatrex/__init__.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: agent.py
Description: Defines an agent, central unit in our simulation

"""
from typing import List
import uuid

from simulatrex.environment import BaseEnvironment
from simulatrex.config import (
    AgentIdentity,
    InitialConditions,
    AgentRelationship,
    AgentGroup,
    Objective,
)
from simulatrex.agent_utils.types import AgentType, AgentMemory, CognitiveModel
from simulatrex.agent_utils.perceive import perceive
from simulatrex.agent_utils.think import think
from simulatrex.event import Event
from simulatrex.utils.logger_config import Logger
from simulatrex.llm_utils.models import OpenAILanguageModel, LlamaLanguageModel
from simulatrex.llm_utils.prompts import PromptManager, TemplateType

logger = Logger()


class Message:
    def __init__(
        self,
        id: str,
        sender_id: str,
        receiver_id: str,
        content: str,
        group_id: str = None,
        metadata: any = {},
    ):
        self.id = id
        self.sender_id = sender_id
        self.receiver_id = receiver_id
        self.content = content
        self.group_id = group_id
        self.metadata = metadata


class BaseAgent:
    def __init__(
        self,
        id: str,
        type: AgentType,
        identity: AgentIdentity,
        initial_conditions: InitialConditions,
    ):
        self.id = "agent_" + id
        self.type = type

        self.identity = identity
        self.initial_conditions = initial_conditions

        self.relationships: List[AgentRelationship] = []
        self.groups: List[AgentGroup] = []

        self.memory = AgentMemory(
            self.id,
            decay_factor=initial_conditions["decay_factor"]
            if "decay_factor" in initial_conditions
            else 0.5,
        )


class LLMAgent(BaseAgent):
    def __init__(
        self,
        id: str,
        type: AgentType,
        identity: AgentIdentity,
        initial_conditions: InitialConditions,
        cognitive_model_id: str,
    ):
        super().__init__(id, type, identity, initial_conditions)

        self.cognitive_model_id = cognitive_model_id
        self.cognitive_model = None
        self.message_queue: List[Message] = []

        logger.debug(
            f"Agent {self.id} is using cognitive model {self.cognitive_model_id}"
        )
        self.init_cognition()

    def init_cognition(self):
        if self.cognitive_model_id == CognitiveModel.GPT_4.value:
            self.cognitive_model = OpenAILanguageModel(
                model_id=CognitiveModel.GPT_4, agent_id=self.id
            )
        elif self.cognitive_model_id == CognitiveModel.LLAMA_2_70b.value:
            self.cognitive_model = LlamaLanguageModel(
                model_id=CognitiveModel.LLAMA_2_70b, agent_id=self.id
            )
        else:
            raise NotImplementedError(
                f"Model {self.cognitive_model_id} not implemented yet"
            )

    async def perceive_event(self, event: Event, environment: BaseEnvironment):
        logger.debug(f"Agent {self.id} is processing event {event.id}")
        await perceive(
            self.cognitive_model,
            self.memory,
            self.identity,
            environment,
            event,
        )

    async def think(self, environment: BaseEnvironment):
        # Based on the provided envrioment make some thoughts
        await think(self.cognitive_model, self.memory, self.identity, environment)

    async def evaluate_outputs(
        self, environment: BaseEnvironment, objective: Objective
    ):
        logger.debug(f"Agent {self.id} is evaluating outputs")
        # Based on agent memory and environment, evaluate outputs
        prompt = PromptManager().get_filled_template(
            TemplateType.AGENT_EVALUATION,
            agent_name=self.identity.name,
            agent_output=self.memory.long_term_memory.retrieve_memory(
                environment.context, n_results=10
            ),
            environment=environment,
            objective=objective,
        )

        response = await self.cognitive_model.ask(prompt)
        return response

    def initiate_conversation(self):
        # Use the cognitive model to decide when to send a message and what the content should be
        if self.cognitive_model.should_send_message():  # You need to implement this
            content = self.cognitive_model.generate_message_content()  # And this
            receiver_id = self.choose_receiver()  # And this
            self.send_message(receiver_id, content)

    def send_message(self, receiver_id: str, content: str):
        message_id = str(uuid.uuid4())
        message = Message(message_id, self.id, receiver_id, content)
        return message

    async def process_messages(self):
        for message in self.message_queue:
            # Use the cognitive model to produce results
            converse_prompt = PromptManager().get_filled_template(
                TemplateType.AGENT_REPLY_TO_CONVERSATION,
                name=self.identity.name,
                age=self.identity.age,
                gender=self.identity.gender,
                persona=self.identity.persona,
                ethnicity=self.identity.ethnicity,
                language=self.identity.language,
                core_memories=", ".join(self.identity.core_memories),
                sender_identity=message.metadata,
                received_message=message.content,
            )

            response_content = await self.cognitive_model.ask(converse_prompt)

            # Send a reply message
            self.send_message(message.sender_id, response_content)

            # Remove current message object from queue
            self.message_queue = [
                msg for msg in self.message_queue if msg.id != message.id
            ]

src/simulatrex/agent.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: converse.py
Description: Converse for inter-agent communication

"""
from simulatrex.environment import BaseEnvironment
from simulatrex.config import AgentIdentity
from simulatrex.llm_utils.prompts import PromptManager, TemplateType
from simulatrex.llm_utils.models import BaseLanguageModel
from simulatrex.utils.logger_config import Logger
from .types import AgentMemory

logger = Logger()


async def converse(
    cognitive_model: BaseLanguageModel,
    memory: tuple[AgentMemory],
    identities: tuple[AgentIdentity],
    environment: BaseEnvironment,
    topic: str,
):
    """
    Converse on a given topic
    """

    # Converse about a current topic
    prompt = PromptManager().get_filled_template(
        TemplateType.AGENT_CONVERSE,
        agent1_name=identities[0].name,
        agent1_age=identities[0].age,
        agent1_gender=identities[0].gender,
        agent1_persona=identities[0].persona,
        agent1_memories=", ".join(identities[0].core_memories),
        agent2_name=identities[1].name,
        agent2_age=identities[1].age,
        agent2_gender=identities[1].gender,
        agent2_persona=identities[1].persona,
        agent2_core_memories=", ".join(identities[1].core_memories),
        environment_context=environment.context,
        topic=topic,
    )

    response = await cognitive_model.ask(prompt)

    logger.debug(response)

    return response

src/simulatrex/agent_utils/converse.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: perceive.py
Description: Perceive the current enviroment and what is happening

"""
from simulatrex.environment import BaseEnvironment
from simulatrex.config import AgentIdentity
from simulatrex.llm_utils.prompts import PromptManager, TemplateType
from simulatrex.llm_utils.models import BaseLanguageModel
from simulatrex.utils.logger_config import Logger
from simulatrex.event import Event
from .types import AgentMemory

logger = Logger()


async def perceive(
    cognitive_model: BaseLanguageModel,
    memory: AgentMemory,
    identity: AgentIdentity,
    environment: BaseEnvironment,
    event: Event,
):
    """
    Perceive the current event and what is happening
    """
    recent_memories = memory.short_term_memory.retrieve_memory(
        event.content, n_results=5
    )

    logger.debug("Recent memories: %s", recent_memories)

    # Based on event content make a decision and request model
    prompt = PromptManager().get_filled_template(
        TemplateType.AGENT_PERCEPTION,
        name=identity.name,
        age=identity.age,
        gender=identity.gender,
        persona=identity.persona,
        ethnicity=identity.ethnicity,
        language=identity.language,
        core_memories=", ".join(identity.core_memories),
        recent_memories=", ".join(recent_memories),
        environment_description=environment.description,
        environment_context=environment.context,
        environment_entities=", ".join(environment.entities),
        event=event,
    )

    response = await cognitive_model.ask(prompt)

    logger.debug(response)

    return response

src/simulatrex/agent_utils/perceive.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: think.py
Description: Think and create memories based  on environment settings

"""
from simulatrex.memory import MemoryUnit
from simulatrex.environment import BaseEnvironment
from simulatrex.config import AgentIdentity
from simulatrex.llm_utils.prompts import PromptManager, TemplateType
from simulatrex.llm_utils.models import BaseLanguageModel
from simulatrex.utils.logger_config import Logger
from .types import AgentMemory

logger = Logger()


async def think(
    cognitive_model: BaseLanguageModel,
    memory: AgentMemory,
    identity: AgentIdentity,
    environment: BaseEnvironment,
):
    """
    Think and create memories based  on environment settings
    """
    prompt = PromptManager().get_filled_template(
        TemplateType.AGENT_THINK,
        name=identity.name,
        age=identity.age,
        gender=identity.gender,
        persona=identity.persona,
        ethnicity=identity.ethnicity,
        language=identity.language,
        core_memories=", ".join(identity.core_memories),
        environment_context=environment.context,
        environment_description=environment.description,
        environment_entities=", ".join(environment.entities),
    )

    response = await cognitive_model.generate_structured_output(prompt, MemoryUnit)

    memory.short_term_memory.add_memory(response)
    memory.long_term_memory.add_memory(response)

    logger.debug(response)

src/simulatrex/agent_utils/think.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: types.py
Description: Shared agent types

"""
from enum import Enum
from typing import Dict, List

from simulatrex.memory import LongTermMemory, ShortTermMemory


class AgentType(Enum):
    LLMAgent = "LLM_AGENT"


class CognitiveModel(Enum):
    GPT_4 = "gpt-4"
    LLAMA_2_70b = "Llama2_70b"


class AgentMemory:
    def __init__(self, id: str, decay_factor: int) -> None:
        self.decay_factor = decay_factor
        self.short_term_memory = ShortTermMemory(id, decay_factor)
        self.long_term_memory = LongTermMemory(id)

src/simulatrex/agent_utils/types.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: config.py
Description: Defines a config for a simulation

"""
from pydantic import BaseModel
from typing import List, Dict, Optional, Union


class Objective(BaseModel):
    id: str
    description: str
    metric: str
    target: str


class ResponseModel(BaseModel):
    type: str
    parameters: Dict[str, float]


class AgentIdentity(BaseModel):
    name: str
    age: int
    gender: str
    ethnicity: str
    language: str
    persona: str
    personality_description: str
    traits: List[str]
    interests: List[str]
    knowledge_base: List[str]
    skills: List[str]
    behavior_patterns: List[str]
    past_experiences: List[str]
    societal_role: str
    affiliations: List[str]
    current_state: str
    core_memories: List[str]


class InitialConditions(BaseModel):
    awareness: float


class AgentRelationship(BaseModel):
    agent_id: str  # ID of the other agent in this relationship
    type: str  # E.g., "friend", "colleague"
    strength: float  # E.g., from 0 (acquaintance) to 1 (best friend)
    metadata: Optional[Dict[str, Union[str, int, float]]]


class AgentGroup(BaseModel):
    id: str
    type: str  # E.g., "school", "company"
    member_agent_ids: List[str]
    metadata: Optional[Dict[str, Union[str, int, float]]]


class Agent(BaseModel):
    id: str
    type: str
    identity: AgentIdentity
    initial_conditions: InitialConditions
    cognitive_model: str
    relationships: List[AgentRelationship]
    group_affiliations: List[str]
    # response_model: ResponseModel


class AgentsHierarchy(BaseModel):
    organizations: List[Dict[str, Union[str, List[str]]]]


class MappingRule(BaseModel):
    csv_column: str
    agent_attribute: str


class DataFeed(BaseModel):
    type: str
    location: str
    mapping_rules: List[MappingRule]


class SimulationTimeConfig(BaseModel):
    start_time: str
    end_time: str
    time_multiplier: int


class Environment(BaseModel):
    type: str
    description: str
    context: str
    entities: List[str]
    time_config: SimulationTimeConfig
    # data_feed: Union[None, DataFeed]


class Event(BaseModel):
    id: str
    type: str
    source: str
    content: str
    impact: float
    scheduled_time: str


class Evaluation(BaseModel):
    metrics: List[str]
    objectives: List[Objective]


class ExpectedOutcome(BaseModel):
    average_awareness_level: float
    highest_influence_platform: str
    agent_collaboration_count: int


class SimulationConfig(BaseModel):
    title: str
    environment: Environment
    agents: List[Agent]
    groups: List[AgentGroup] or None
    events: List[Event]
    evaluation: Evaluation or None
    # expected_outcomes: Optional[ExpectedOutcome]


class Config(BaseModel):
    version: str
    simulation: SimulationConfig

src/simulatrex/config.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: engine.py
Description: Main engine for running simulations

"""
from typing import List

from simulatrex.config import Config
from simulatrex.agent import LLMAgent
from simulatrex.environment import (
    StaticEnvironment,
    DynamicEnvironment,
    EnvironmentType,
)
from simulatrex.evaluation import EvaluationEngine
from simulatrex.utils.json_utils import JSONHelper
from simulatrex.utils.logger_config import Logger


class SimulationEngine:
    def __init__(self, config_path, num_iterations=None):
        json_data = JSONHelper.read_json(config_path)
        self.config = Config(**json_data)
        self.title = self.config.simulation.title

        self.agents = self.init_agents()  # Initialize agents from JSON
        self.environment = self.init_environment()  # Initialize environment from JSON

        time_config = self.config.simulation.environment.time_config
        self.environment.init_time(time_config)

        self.current_iteration = 1

    def init_agents(self) -> List[LLMAgent]:
        agents = []

        for agent in self.config.simulation.agents:
            new_agent = LLMAgent(
                agent.id,
                agent.type,
                agent.identity,
                agent.initial_conditions,
                agent.cognitive_model,
            )

            agents.append(new_agent)

        return agents

    def init_environment(self):
        environment_settings = self.config.simulation.environment
        enviroment_type = environment_settings.type
        events = self.config.simulation.events

        if enviroment_type == EnvironmentType.Static.value:
            new_environment = StaticEnvironment(
                environment_settings.description,
                environment_settings.context,
                environment_settings.entities,
                events,
            )
            return new_environment
        elif enviroment_type == EnvironmentType.Dynamic.value:
            new_environment = DynamicEnvironment(
                environment_settings.description,
                environment_settings.context,
                environment_settings.entities,
                events,
            )
            return new_environment
        else:
            raise ValueError(f"Unsupported environment type: {enviroment_type}")

    async def run(self):
        # Append the title to the logfiles to distinguish between different simulations
        logger = Logger(
            name=self.title,
            log_file=f"{self.title}_run.log",
            response_log_file=f"{self.title}_agents_response.log",
        )

        while True:
            # Check stopping time
            if not self.environment.is_running():
                break

            # Update the environment
            recent_events, current_env_context = await self.environment.update()

            # Log the current env context
            logger.debug(f"Current environment context: {current_env_context}")

            # Let the agents process the recent events
            for agent in self.agents:
                # Agent thinks about environment context
                await agent.think(self.environment)

                # Agent perceives the recent events
                for event in recent_events:
                    logger.debug(f"Event - ID: {event.id}, Content: {event.content}")
                    await agent.perceive_event(event, self.environment)
                    await agent.process_messages()

            # Log the current iteration
            logger.info(
                f"Simulation Iteration {self.current_iteration}: Processed recent events."
            )

            self.current_iteration += 1

        # Evaluate the simulation
        evaluation_engine = EvaluationEngine(self.config.simulation.evaluation)
        evaluation_engine.evaluate_agents_outputs(self.agents, self.environment)

src/simulatrex/engine.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: environment.py
Description: Defines the environment class, with a base environment and a dynamic and static environment

"""
from datetime import timedelta
from enum import Enum
from typing import List

from simulatrex.agent_utils.types import CognitiveModel
from simulatrex.llm_utils.models import OpenAILanguageModel
from simulatrex.llm_utils.prompts import PromptManager, TemplateType
from simulatrex.utils.time_utils import TimeUtils
from simulatrex.event import Event, EventEngine
from simulatrex.config import SimulationTimeConfig


class EnvironmentType(Enum):
    Dynamic = "DYNAMIC"
    Static = "STATIC"


class BaseEnvironment:
    def __init__(
        self, description: str, context: str, entities: List[str], events: List[any]
    ):
        self.description = description
        self.context = context  # Contextual information about the environment
        self.entities = entities

        self.event_engine = self.init_event_engine(
            events
        )  # Initialize events from JSON

        self.current_time = 0
        self.start_time = 0
        self.end_time = 0
        self.type = EnvironmentType.Static

        self.llm = OpenAILanguageModel(model_id=CognitiveModel.GPT_4)

    def init_event_engine(self, events: List[any]):
        event_engine = EventEngine(events)
        return event_engine

    def init_time(self, time_config: SimulationTimeConfig):
        self.start_time = TimeUtils.to_timestamp(time_config.start_time)
        self.end_time = TimeUtils.to_timestamp(time_config.end_time)

        self.time_multiplier = time_config.time_multiplier
        self.current_time = self.start_time

    def get_current_time(self):
        return self.current_time

    def is_running(self):
        return self.current_time < self.end_time

    async def update(self) -> tuple[List[Event], str]:
        # Store the previous time before updating the current_time
        previous_time = self.current_time

        # Increase time
        time_step = timedelta(seconds=1).total_seconds() * self.time_multiplier
        self.current_time += time_step

        # Process events
        active_events = self.event_engine.process_events(
            previous_time, self.current_time
        )

        # Summarize the current context
        current_env_context = await self.summarize_context(active_events)

        # Update current env context
        self.context = current_env_context

        return (active_events, current_env_context)

    async def summarize_context(self, active_events: List[Event]):
        # Process current events and summarize the context
        summarize_env_prompt = PromptManager().get_filled_template(
            TemplateType.SUMMARIZE_ENVIRONMENT,
            environment_type=self.type.value,
            environment_context=self.context,
            environment_entities=", ".join(self.entities),
            environment_description=self.description,
            events_summary=", ".join([event.content for event in active_events]),
        )

        response = await self.llm.ask(summarize_env_prompt)
        return response

    def get_recent_events(self, max_events=5):
        return self.event_engine.get_recent_events(max_events)


class StaticEnvironment(BaseEnvironment):
    def __init__(
        self, description: str, context: str, entities: List[str], events: List[any]
    ):
        super().__init__(description, context, entities, events)
        self.type = EnvironmentType.Static


class DynamicEnvironment(BaseEnvironment):
    def __init__(
        self, description: str, context: str, entities: List[str], events: List[any]
    ):
        super().__init__(description, context, entities, events)
        self.type = EnvironmentType.Dynamic

src/simulatrex/environment.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: evaluation.py
Description: Evaluation Processor

"""
from typing import List

from simulatrex.agent import LLMAgent
from simulatrex.config import Evaluation
from simulatrex.environment import BaseEnvironment
from simulatrex.llm_utils.prompts import PromptManager, TemplateType
from simulatrex.llm_utils.models import OpenAILanguageModel
from simulatrex.agent_utils.types import CognitiveModel


class Objective:
    def __init__(self, id: str, description: str, metric: str, target: str):
        self.id = id
        self.description = description
        self.metric = metric
        self.target = target


class EvaluationEngine:
    def __init__(self, evaluation: Evaluation):
        self.metrics = evaluation.metrics
        self.objectives = [Objective(**obj) for obj in evaluation.objectives]

        self.llm = OpenAILanguageModel(model_id=CognitiveModel.GPT_4)
        self.prompt_manager = PromptManager()

    async def evaluate_agents_outputs(
        self,
        agents: List[LLMAgent],
        environment: BaseEnvironment,
    ):
        agent_results = ""
        results = []

        for objective in self.objectives:
            for agent in agents:
                agent_result = await agent.evaluate_outputs(environment, objective)
                agent_results += agent_result + "\n\n"

            prompt = self.prompt_manager.get_filled_template(
                TemplateType.EVALUATE_AGENT_OUTPUTS,
                agent_results=agent_results,
                environment=environment,
                objective=objective,
            )

            llm_response = await self.llm.ask(prompt)

            results.append(
                {
                    "objective_id": objective.id,
                    "description": objective.description,
                    "llm_response": llm_response,
                }
            )

        return results

src/simulatrex/evaluation.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: event.py
Description: Defines an event for a simulation and the event engine to manage events.

"""
from typing import Any, List, Union

from simulatrex.utils.logger_config import Logger
from simulatrex.utils.time_utils import TimeUtils


class Event:
    def __init__(
        self,
        id: str,
        type: str,
        source: str = None,
        content: str = None,
        impact: float = 0.0,
        scheduled_time: str = None,
    ):
        """
        Initialize an event.

        Args:
            id (str): Unique identifier for the event.
            type (str): Type of the event.
            source (str, optional): Source of the event. Defaults to None.
            content (str, optional): content of the event. Defaults to None.
            impact (float, optional): Impact of the event. Defaults to 0.0.
            scheduled_time (str, optional): Scheduled time of the event. Defaults to None.
        """
        self.id = id
        self.type = type
        self.source = source
        self.content = content
        self.impact = impact
        self.scheduled_time = TimeUtils.to_timestamp(scheduled_time)
        self.is_triggered = False

    def __str__(self):
        return f"Event(id={self.id}, type={self.type}, source={self.source}, content={self.content}, scheduled_time={self.scheduled_time})"


class EventEngine:
    def __init__(self, events: List[any]):
        self.events = [
            Event(
                event.id,
                event.type,
                event.source,
                event.content,
                event.impact,
                event.scheduled_time,
            )
            for event in events
        ]

        self.init_events()

    def init_events(self):
        # Clear the list of events
        self.active_events: List[Event] = []
        self.past_events: List[Event] = []

    def process_events(self, previous_time: int, current_time: int):
        # Iterate over all events and check their conditions
        self.init_events()
        for event in self.events:
            if self._should_trigger(event, previous_time, current_time):
                self.active_events.append(event)
                event.is_triggered = True
            else:
                self.past_events.append(event)

        return self.active_events

    def _should_trigger(self, event: Event, previous_time: int, current_time: int):
        # Placeholder method. This should check the event's trigger conditions
        # and return True if they are met.
        # For simplicity, let's check against a time condition here.
        return (
            previous_time <= event.scheduled_time < current_time
            and not event.is_triggered
        )

    def get_recent_events(self, max_events=5):
        return self.active_events[-max_events:]

src/simulatrex/event.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: models.py
Description: LLM Models

"""
import os

from dotenv import load_dotenv
from pydantic import BaseModel
import openai
import requests
import instructor

load_dotenv()

from simulatrex.agent_utils.types import CognitiveModel
from simulatrex.utils.logger_config import Logger

logger = Logger()

DEFAULT_SYSTEM_PROMPT = "This is a real-world like simulation. Try to answer the following questions as best as possible:"


class BaseLanguageModel:
    async def ask(
        self, prompt: str, context_prompt=DEFAULT_SYSTEM_PROMPT, temperature=0.9
    ) -> str:
        pass

    async def generate_structured_output(
        self,
        prompt: str,
        response_model: BaseModel,
        context_prompt=DEFAULT_SYSTEM_PROMPT,
        temperature=0.9,
    ) -> dict:
        pass


class OpenAILanguageModel(BaseLanguageModel):
    """
    This is a wrapper for the OpenAI API.
    """

    def __init__(self, model_id=CognitiveModel.GPT_4, agent_id=None):
        api_key = os.environ.get("OPENAI_API_KEY")

        if api_key is None:
            raise Exception(
                "No OpenAI API key found. Please set OPENAI_API_KEY as environment variable."
            )

        self.api_key = api_key
        self.model_id = model_id
        self.agent_id = agent_id

        # Patch the openai.ChatCompletion.create and openai.ChatCompletion.acreate methods to support the response_model parameter.
        instructor.patch()

    async def ask(
        self, prompt: str, context_prompt=DEFAULT_SYSTEM_PROMPT, temperature=0.9
    ) -> str:
        try:
            chat_completion = openai.ChatCompletion.create(
                messages=[
                    {"role": "system", "content": context_prompt},
                    {"role": "user", "content": prompt},
                ],
                model=self.model_id.value,
                max_retries=2,
                temperature=temperature,
            )
            response_message = chat_completion["choices"][0].message.content

            # Log the response
            if self.agent_id:
                logger.log_agent_response(self.agent_id, response_message)

            return response_message
        except Exception as e:
            logger.debug(f"Error: {e}")

            raise Exception("Error in OpenAI API call")

    async def generate_structured_output(
        self,
        prompt: str,
        response_model: BaseModel,
        context_prompt=DEFAULT_SYSTEM_PROMPT,
        temperature=0.9,
    ):
        try:
            structured_response = openai.ChatCompletion.create(
                messages=[
                    {"role": "system", "content": context_prompt},
                    {"role": "user", "content": prompt},
                ],
                response_model=response_model,
                model=self.model_id.value,
                max_retries=2,
                temperature=temperature,
            )
            return structured_response
        except Exception as e:
            logger.debug(f"Error: {e}")

            raise Exception("Error in OpenAI API call")


class LlamaLanguageModel(BaseLanguageModel):
    """
    This class is a wrapper for the LLama API.
    """

    def __init__(self, model_id=CognitiveModel.LLAMA_2_70b, agent_id=None):
        access_token = os.environ.get("HUGGINGFACE_ACCESS_TOKEN")

        if access_token is None:
            raise Exception(
                "No Huggingface API key found. Please set HUGGINGFACE_API_KEY as environment variable."
            )

        self.access_token = access_token
        self.model_id = model_id
        self.agent_id = agent_id

    async def ask(self, prompt: str) -> str:
        if self.model_id == CognitiveModel.LLAMA_2_70b:
            try:
                API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-2-70b-chat-hf"
                headers = {"Authorization": f"Bearer {self.access_token}"}

                response = requests.post(
                    API_URL,
                    headers=headers,
                    json={
                        "inputs": prompt,
                        "options": {
                            "use_cache": False,
                            "wait_for_model": True,
                        },
                        "parameters": {
                            "max_length": 4096,
                        },
                    },
                )

                result = response.json()
                result_text = result["generated_text"]

                # Log the response
                if self.agent_id:
                    logger.log_agent_response(self.agent_id, result_text)

                return result_text
            except Exception as e:
                logger.debug(f"Error: {e}")

                raise Exception("Error in HuggingFace API call")

src/simulatrex/llm_utils/models.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: prompts.py
Description: Prompt Templates
Prompt Templates are used to generate prompts for the LLM.

"""
from enum import Enum
import os


class TemplateType(Enum):
    AGENT_PERCEPTION = "agent_perception.txt"
    AGENT_CONVERSE = "agent_converse.txt"
    AGENT_THINK = "agent_think.txt"
    AGENT_EVALUATION = "agent_evaluation.txt"
    AGENT_START_CONVERSATION = "agent_start_conversation.txt"
    AGENT_REPLY_TO_CONVERSATION = "agent_reply_to_conversation.txt"
    SUMMARIZE_ENVIRONMENT = "summarize_environment.txt"
    EVALUATE_AGENT_OUTPUTS = "evaluate_agent_outputs.txt"
    AGENT_INTERVIEW = "agent_interview.txt"


class PromptManager:
    def __init__(self, template_dir="prompt_templates"):
        self.template_dir = template_dir
        self.templates = self._load_templates()

    def _load_templates(self):
        templates = {}
        for template_type in TemplateType:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            template_dir_path = os.path.join(current_dir, self.template_dir)
            with open(os.path.join(template_dir_path, template_type.value), "r") as f:
                templates[template_type] = f.read()
        return templates

    def get_filled_template(self, template_type: TemplateType, **attributes) -> str:
        if template_type not in self.templates:
            raise ValueError(f"Template type {template_type} not found!")

        return self.templates[template_type].format(**attributes)

src/simulatrex/llm_utils/prompts.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: memory.py
Description: Defines the memory class, which is used to define the memory of an agent

"""
import math
import uuid
from typing import List

from pydantic import BaseModel

from simulatrex.utils.time_utils import TimeUtils
from simulatrex.vectordb import VectorDB


class MemoryUnit(BaseModel):
    type: str
    depth: float
    content: str
    keywords: List[str]

    id: str = str(uuid.uuid4())
    created: int = 0
    last_accessed: int = 0

    def get_metadata(self):
        return {
            "id": self.id,
            "type": self.type,
            "depth": self.depth,
            "created": self.created,
            "last_accessed": self.last_accessed,
            "keywords": ",".join(self.keywords),
        }


class BaseMemory:
    def __init__(self, id: str):
        self.id = id
        self.vector_db = VectorDB(collection_name=id)

    def add_memory(self, memory_unit: MemoryUnit):
        self.vector_db.add_memory(
            memory_unit.content, metadatas=[memory_unit.get_metadata()]
        )

    def retrieve_memory(self, content: str, n_results: int):
        results = self.vector_db.query_memory(content, n_results=n_results)
        return results


class ShortTermMemory(BaseMemory):
    def __init__(self, id: str, decay_factor: int):
        stm_id = "stm_" + id
        super().__init__(stm_id)
        self.decay_factor = decay_factor

    def retrieve_memory(self, content: str, n_results: int, current_timestamp: int):
        results = self.vector_db.query_memory(content, n_results=n_results)

        # Apply decay factor to the results based on the 'created' date
        for memory in results:
            time_diff = current_timestamp - TimeUtils.to_timestamp(
                memory["last_accessed"]
            )
            memory["score"] *= math.exp(-self.decay_factor * time_diff)

        # Sort the results by score in descending order
        results.sort(key=lambda x: x["score"], reverse=True)
        return results


class LongTermMemory(BaseMemory):
    def __init__(self, id: str):
        ltm_id = "ltm_" + id
        super().__init__(ltm_id)

src/simulatrex/memory.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: json_utils.py
Description: Helpers for json

"""
import json
from typing import Any, Dict


class JSONHelper:
    @staticmethod
    def read_json(file_path: str) -> Dict[str, Any]:
        with open(file_path, "r", encoding="utf-8") as f:
            return json.load(f)

    @staticmethod
    def write_json(file_path: str, data: Dict[str, Any]) -> None:
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=4)

src/simulatrex/utils/json_utils.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: log_config.py
Description: Log helper

"""
import logging
import json

from termcolor import colored


class Logger:
    def __init__(
        self,
        name="simulatrex-agents",
        log_file="run.log",
        response_log_file="response.log",
    ):
        # Create logger
        self.logger = logging.getLogger(name)
        self.agent_logger = logging.getLogger("agent_logger")
        if not self.logger.handlers:  # Check if logger already has handlers
            self.logger.setLevel(logging.DEBUG)  # Set default logging level

            # Create console handler and set level to debug
            ch = logging.StreamHandler()
            ch.setLevel(logging.DEBUG)

            # Create file handler and set level to debug
            fh = logging.FileHandler(log_file)
            fh.setLevel(logging.DEBUG)

            # Create file handler for response log
            rh = logging.FileHandler(response_log_file)
            rh.setLevel(logging.INFO)

            # Create formatter
            formatter = logging.Formatter(
                "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
            )

            # Add formatter to handlers
            ch.setFormatter(formatter)
            fh.setFormatter(formatter)
            rh.setFormatter(formatter)

            # Add handlers to logger
            self.logger.addHandler(ch)
            self.logger.addHandler(fh)

            self.agent_logger.addHandler(rh)

    def debug(self, msg):
        self.logger.debug(colored(msg, "blue"))

    def info(self, msg):
        self.logger.info(colored(msg, "green"))

    def warning(self, msg):
        self.logger.warning(colored(msg, "yellow"))

    def error(self, msg):
        self.logger.error(colored(msg, "red"))

    def log_agent_response(self, agent_id: str, response: str):
        response_dict = {"agent_id": agent_id, "response": response}
        self.agent_logger.info(json.dumps(response_dict), extra={"handler": "rh"})

src/simulatrex/utils/logger_config.py
simulatrex/simulatrex
from datetime import datetime


class TimeUtils:
    @staticmethod
    def to_timestamp(datetime_str: str) -> int:
        """Converts a datetime string to a timestamp."""
        dt = datetime.fromisoformat(datetime_str)
        return int(dt.timestamp())

    @staticmethod
    def from_timestamp(timestamp: int) -> str:
        """Converts a timestamp back to a datetime string."""
        dt = datetime.fromtimestamp(timestamp)
        return dt.isoformat()

src/simulatrex/utils/time_utils.py
simulatrex/simulatrex
"""
Author: Dominik Scherm (dom@simulatrex.ai)

File: vectordb.py
Description: Defines the vectordb class, which is a wrapper around ChromaDB

"""
import os
import chromadb
from chromadb.utils import embedding_functions
from dotenv import load_dotenv

load_dotenv()

from simulatrex.utils.logger_config import Logger

logger = Logger()


class VectorDB:
    def __init__(self, collection_name: str):
        api_key = os.getenv("OPENAI_API_KEY")

        if api_key is None:
            raise Exception(
                "No OpenAI API key found. Please set OPENAI_API_KEY as environment variable."
            )

        if collection_name is None:
            raise Exception("No collection name provided.")

        self.collection_name = collection_name
        self.api_key = api_key
        self.init_chroma()

    def init_chroma(self):
        try:
            # Embed using openai, TODO: replace with abstract embedding
            openai_embed_function = embedding_functions.OpenAIEmbeddingFunction(
                api_key=self.api_key,
                # TODO: replace with abstract model name
                model_name="text-embedding-ada-002",
            )

            current_dir = os.path.dirname(os.path.abspath(__file__))
            chroma_db_path = os.path.join(current_dir, "chromadb")
            self.client = chromadb.PersistentClient(path=chroma_db_path)

            self.collection = self.client.create_collection(
                name=self.collection_name,
                embedding_function=openai_embed_function,
                get_or_create=True,
            )

        except Exception as e:
            logger.error(e)
            raise Exception("Failed to initialize ChromaDB collection.")

    def add_memory(self, content: str, metadatas=None, ids=None):
        """
        Add memory which will be vectorized

        Parameters:
        - content: content string
        - metadatas (list, optional): List of metadata dictionaries for each document. Defaults to None.
        - ids (list, optional): List of unique IDs for each document. Defaults to None.
        """

        self.collection.add(
            documents=[content],
            metadatas=metadatas or [{}],  # Empty metadata if not provided
            ids=ids or [f"doc.1"],  # Default ID generation if not provided
        )

    def query_memory(
        self, query_text, n_results=5, where_metadata=None, where_document=None
    ):
        """
        Query the ChromaDB collection for semantically similar documents.

        Parameters:
        - query_text (str): The text to query against the stored documents.
        - n_results (int, optional): Number of results to return. Defaults to 5.
        - where_metadata (dict, optional): Filters based on metadata fields. Defaults to None.
        - where_document (dict, optional): Filters based on document content. Defaults to None.

        Returns:
        - list: List of query results.
        """
        results = self.collection.query(
            query_texts=[query_text],
            n_results=n_results,
            where=where_metadata,
            where_document=where_document,
        )
        return results

    def delete_memory(self, memory_id: str):
        """
        Delete a memory by id

        Parameters:
        - id (str): The id of the memory to delete
        """
        self.collection.delete([memory_id])

    def get_collection(self, collection_id: str):
        """
        Get a collection by id

        Parameters:
        - id (str): The id of the memory to get
        """
        return self.client.get_collection(collection_id)

    def delete_collection(self, collection_id: str):
        """
        Delete a collection by id

        Parameters:
        - id (str): The id of the memory to delete
        """
        self.client.delete_collection(collection_id)

src/simulatrex/vectordb.py
simulatrex/simulatrex