feature(example): add a new demo script and fix some bugs (#7)

* add run_chat and fix some bugs

* delete some chinese notes

Author: SwimmingWang

examples/run_chat.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+"""
+Simple script to run a multi-agent chat conversation
+Usage: python scripts/run_chat.py
+"""
+
+import asyncio
+import os
+import sys
+from pathlib import Path
+
+# Add the project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from tiny_chat.server import ChatServer
+from tiny_chat.messages import ChatBackground
+
+
+async def main():
+    """Run a simple multi-agent conversation"""
+    
+    # Get API key from environment variable
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        print("Warning: OPENAI_API_KEY not set. Some features may not work.")
+    
+    # Create chat server
+    server = ChatServer(api_key=api_key)
+    
+    # Define agent configurations
+    agent_configs = [
+        {
+            "name": "Alice",
+            "type": "llm",
+            "model": "gpt-4o-mini",
+            "goal": "Be friendly and helpful in the conversation"
+        },
+        {
+            "name": "Bob", 
+            "type": "llm",
+            "model": "gpt-4o-mini",
+            "goal": "Ask thoughtful questions and share interesting ideas"
+        }
+    ]
+    
+    # Create a simple background
+    background = ChatBackground(
+        scenario="Two friends meeting at a coffee shop",
+        p1_background="Alice is a software engineer who loves hiking",
+        p2_background="Bob is a teacher who enjoys reading science fiction",
+        p1_goal="Have a pleasant conversation about weekend plans",
+        p2_goal="Discuss recent books and outdoor activities",
+        p1_name="Alice",
+        p2_name="Bob"
+    )
+    
+    print("Starting multi-agent conversation...")
+    print("=" * 50)
+    
+    # Run the conversation
+    await server.run_conversation(
+        agent_configs=agent_configs,
+        background=background,
+        max_turns=10,
+        enable_evaluation=True
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main()) 

tiny_chat/__init__.py

@@ -5,20 +5,19 @@
 __version__ = '0.0.1'
 
 from .agents import LLMAgent
-from .environment import ChatEnvironment
-from .evaluators import (
+from .envs import TinyChatEnvironment
+from .evaluator import (
     EpisodeLLMEvaluator,
     EvaluationDimension,
     EvaluationForMultipleAgents,
     EvaluationForTwoAgents,
     Evaluator,
     RuleBasedTerminatedEvaluator,
-    SotopiaDimensions,
+    TinyChatDimensions,
     unweighted_aggregate_evaluate,
 )
-from .generator import MessageGenerator
 from .messages import AgentAction, ChatBackground, Message, Observation, SimpleMessage
-from .profile import AgentProfile, RelationshipProfile, RelationshipType
+from .profile import BaseAgentProfile, BaseEnvironmentProfile, BaseRelationshipProfile
 from .server import ChatServer
 
 __all__ = [
@@ -28,15 +27,14 @@
     'Observation',
     'AgentAction',
     'ChatBackground',
-    'ChatEnvironment',
-    'MessageGenerator',
-    'AgentProfile',
-    'RelationshipProfile',
-    'RelationshipType',
+    'TinyChatEnvironment',
+    'BaseAgentProfile',
+    'BaseEnvironmentProfile',
+    'BaseRelationshipProfile',
     'Evaluator',
     'RuleBasedTerminatedEvaluator',
     'EpisodeLLMEvaluator',
-    'SotopiaDimensions',
+    'TinyChatDimensions',
     'EvaluationDimension',
     'EvaluationForTwoAgents',
     'EvaluationForMultipleAgents',

tiny_chat/agents/__init__.py

@@ -1,3 +1,6 @@
-from .agents import LLMAgent
+from .agents import LLMAgent, HumanAgent
 
-__all__ = ['LLMAgent']
+__all__ = [
+    'LLMAgent',
+    'HumanAgent',
+]

tiny_chat/agents/agents.py

@@ -1,6 +1,6 @@
-from tiny_chat.generator import generate_agent
+from tiny_chat.generator import agenerate_action
 from tiny_chat.messages import AgentAction, Message, Observation
-from tiny_chat.profile import AgentProfile
+from tiny_chat.profile import BaseAgentProfile
 
 
 class LLMAgent:
@@ -10,17 +10,20 @@ def __init__(
         self,
         agent_name: str,
         model: str = 'gpt-4o-mini',
+        agent_number: int = 1,
         api_key: str | None = None,
         goal: str | None = None,
-        agent_profile: AgentProfile | None = None,
+        agent_profile: BaseAgentProfile | None = None,
+        script_like: bool = False,
     ):
         self.agent_name = agent_name
         self.model = model
+        self.agent_number = agent_number
         self.api_key = api_key
         self._goal = goal or 'Have a natural conversation'
         self.agent_profile = agent_profile
         self.message_history: list[Message] = []
-
+        self.script_like = script_like
     @property
     def goal(self) -> str:
         """Get the agent's goal"""
@@ -32,12 +35,20 @@ def goal(self, goal: str) -> None:
         self._goal = goal
 
     async def act(self, obs: Observation) -> AgentAction:
-        self.recv_message('Environment', obs)
+        self.receive_message('Environment', obs)
 
         if len(obs.available_actions) == 1 and 'none' in obs.available_actions:
             return AgentAction(action_type='none', argument='')
         else:
-            action = await generate_agent()
+            action = await agenerate_action(
+                model_name=self.model,
+                history=self._build_context(obs),
+                turn_number=obs.turn_number,
+                action_types=obs.available_actions,
+                agent=self.agent_name,
+                goal=self.goal,
+                script_like=self.script_like,
+            )
             return action
 
     def _build_context(self, observation: Observation) -> str:
@@ -67,6 +78,9 @@ def reset(self) -> None:
         """Reset agent state"""
         self.message_history.clear()
 
-    def receive_message(self, message: Message) -> None:
+    def receive_message(self, source: str, message: Message) -> None:
         """Receive a message and add to history"""
         self.message_history.append(message)
+
+class HumanAgent:
+    pass

tiny_chat/envs/environment.py

@@ -6,7 +6,7 @@
 
 from pydantic import validate_call
 
-from tiny_chat.evaluators import Evaluator, unweighted_aggregate_evaluate
+from tiny_chat.evaluator import Evaluator, unweighted_aggregate_evaluate
 from tiny_chat.messages import (
     ActionType,
     AgentAction,
@@ -211,6 +211,45 @@ def reset(
                 ),
             ),
         }
+    def get_turn_number(self) -> int:
+        return self.turn_number
+    
+    def is_terminated(self) -> bool:
+        return self.turn_number >= self.max_turns
+    
+    def get_observation(self, agent_name: str) -> Observation:
+        # get last turn
+        last_turn = ""
+        if self.turn_number > 0 and self.inbox:
+            last_actions = {}
+            for source, message in self.inbox:
+                if isinstance(message, AgentAction) and source != 'Environment':
+                    last_actions[source] = message
+            if last_actions:
+                last_turn = _actions_to_natural_language(last_actions)
+            else:
+                last_turn = self.background.to_natural_language()
+        else:
+            last_turn = self.background.to_natural_language()
+        
+        agent_index = 0
+        if hasattr(self, 'agents') and self.agents:
+            try:
+                agent_index = self.agents.index(agent_name)
+            except ValueError:
+                agent_index = 0
+        
+        available_actions = ['none']
+        if hasattr(self, 'action_mask') and len(self.action_mask) > agent_index:
+            if self.action_mask[agent_index]:
+                available_actions = list(self.available_action_types)
+        
+        obs = Observation(
+            last_turn=last_turn,
+            turn_number=self.turn_number,
+            available_actions=available_actions,
+        )
+        return obs
 
     def _get_agent_background(self, agent: Any, agent_id: int) -> str:
         """Extract background information from agent"""

tiny_chat/evaluator/__init__.py

@@ -0,0 +1,21 @@
+from .evaluators import (
+    EpisodeLLMEvaluator,
+    EvaluationDimension,
+    EvaluationForMultipleAgents,
+    EvaluationForTwoAgents,
+    Evaluator,
+    RuleBasedTerminatedEvaluator,
+    TinyChatDimensions,
+    unweighted_aggregate_evaluate,
+)
+
+__all__ = [
+    'Evaluator',
+    'EpisodeLLMEvaluator',
+    'EvaluationDimension',
+    'EvaluationForMultipleAgents',
+    'EvaluationForTwoAgents',
+    'RuleBasedTerminatedEvaluator',
+    'TinyChatDimensions',
+    'unweighted_aggregate_evaluate',
+]

tiny_chat/evaluator/evaluators.py

@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel, validate_call
 
-from .messages import (
+from tiny_chat.messages import (
     AgentAction,
     Message,
     ScriptEnvironmentResponse,
@@ -98,11 +98,11 @@ class EpisodeLLMEvaluator(Evaluator, Generic[T_eval_dim]):
     def __init__(
         self,
         model_name: str,
-        response_format_class: type[EvaluationForTwoAgents[T_eval_dim]],
+        # response_format_class: type[EvaluationForTwoAgents[T_eval_dim]],
     ) -> None:
         self.model_name = model_name
         self.prompt = ''
-        self.response_format_class = response_format_class
+        # self.response_format_class = response_format_class
 
     def __call__(
         self, turn_number: int, messages: list[tuple[str, Message]]
@@ -155,7 +155,7 @@ class EvaluationDimension(BaseModel):
     pass
 
 
-class SotopiaDimensions(BaseModel):
+class TinyChatDimensions(BaseModel):
     """Evaluation dimensions used in Sotopia"""
 
     overall_score: tuple[str, float] = ('Overall score', 0.0)

tiny_chat/generator/__init__.py

@@ -0,0 +1,18 @@
+from .generate_template import agenerate_action, agenerate_goal
+
+from .output_parsers import (
+    EnvResponse,
+    StrOutputParser,
+    ScriptOutputParser,
+    PydanticOutputParser,
+    ListOfIntOutputParser,
+)
+__all__ = [
+    "EnvResponse",
+    "StrOutputParser",
+    "ScriptOutputParser",
+    "PydanticOutputParser",
+    "ListOfIntOutputParser",
+    "agenerate_action",
+    "agenerate_goal",
+]

tiny_chat/generator/generate_template.py

@@ -11,22 +11,22 @@
 from pydantic import validate_call
 from rich import print
 from rich.logging import RichHandler
-from sotopia.generation_utils.output_parsers import (
+from tiny_chat.generator.output_parsers import (
     EnvResponse,
     OutputParser,
     OutputType,
     PydanticOutputParser,
     ScriptOutputParser,
     StrOutputParser,
 )
-from sotopia.utils import format_docstring
+from tiny_chat.utils import format_docstring
 
 from tiny_chat.messages import ActionType, AgentAction, ScriptBackground
-from tiny_chat.messages.message_classes import (
+from tiny_chat.messages import (
     ScriptInteraction,
     ScriptInteractionReturnType,
 )
-from tiny_chat.profile import EnvironmentProfile, RelationshipProfile
+from tiny_chat.profile import BaseEnvironmentProfile, BaseRelationshipProfile
 
 # Configure logger
 log = logging.getLogger('sotopia.generation')
@@ -187,7 +187,7 @@ async def agenerate_env_profile(
     temperature: float = 0.7,
     bad_output_process_model: str | None = None,
     use_fixed_model_version: bool = True,
-) -> EnvironmentProfile:
+) -> BaseEnvironmentProfile:
     """
     Using langchain to generate the background
     """
@@ -204,7 +204,7 @@ async def agenerate_env_profile(
             inspiration_prompt=inspiration_prompt,
             examples=examples,
         ),
-        output_parser=PydanticOutputParser(pydantic_object=EnvironmentProfile),
+        output_parser=PydanticOutputParser(pydantic_object=BaseEnvironmentProfile),
         temperature=temperature,
         bad_output_process_model=bad_output_process_model,
         use_fixed_model_version=use_fixed_model_version,
@@ -217,7 +217,7 @@ async def agenerate_relationship_profile(
     agents_profiles: list[str],
     bad_output_process_model: str | None = None,
     use_fixed_model_version: bool = True,
-) -> tuple[RelationshipProfile, str]:
+) -> tuple[BaseRelationshipProfile, str]:
     """
     Using langchain to generate the background
     """
@@ -232,7 +232,7 @@ async def agenerate_relationship_profile(
         input_values=dict(
             agent_profile=agent_profile,
         ),
-        output_parser=PydanticOutputParser(pydantic_object=RelationshipProfile),
+        output_parser=PydanticOutputParser(pydantic_object=BaseRelationshipProfile),
         bad_output_process_model=bad_output_process_model,
         use_fixed_model_version=use_fixed_model_version,
     )