Fix processing LLM response (#324)

* process dict response

* fix

* fix stream state

Author: braisedpork1964

lagent/actions/tmux_action.py

@@ -112,7 +112,7 @@ def _start(self) -> None:
             capture_output=True,
             text=True,
         )
-        if result.returncode != 0 and result.stderr != 'duplicate session: terminus2\n':
+        if result.returncode != 0 and result.stderr != f'duplicate session: {self._session_name}\n':
             raise RuntimeError(f"Failed to start tmux session: {result.stderr!r}")
         subprocess.run(
             f"tmux set-option -g history-limit {self._history_limit}",

lagent/actions/web_visitor.py

@@ -173,7 +173,11 @@ async def _read_webpage(self, url: str, goal: str) -> str:
                     llm_response = (
                         llm_response.content
                         if isinstance(llm_response, AgentMessage)
-                        else llm_response.choices[0].message.content
+                        else (
+                            llm_response['choices'][0]['message']['content']
+                            if isinstance(llm_response, dict)
+                            else llm_response.choices[0].message.content
+                        )
                     )
                 if not llm_response or len(llm_response) < 10:
                     tool_response = tool_response[: int(len(tool_response) * 0.7)]

lagent/llms/model.py

@@ -2,13 +2,13 @@
 import json
 import random
 import traceback
-
 from typing import Dict, List, Optional, TypedDict, Union
 
 import aiohttp
 
 from lagent.llms.openai import AsyncGPTAPI
 from lagent.utils import get_logger
+
 logger = get_logger()
 
 
@@ -54,7 +54,7 @@ def __init__(
         self.max_tool_response_length = max_tool_response_length
         self.max_tool_calls_per_turn = max_tool_calls_per_turn
 
-    async def chat(self, messages: List[dict], tools=None, **gen_params) -> str:
+    async def chat(self, messages: List[dict], tools=None, **gen_params) -> dict:
         """Generate completion from a list of templates.
 
         Args:

lagent/schema.py

@@ -136,7 +136,8 @@ def from_model_response(cls, response: Union[ChatCompletion, dict], sender: str)
             raw_content_logprobs=msg.get('raw_content_logprobs'),
             extra_info=msg.get('extra_info') or {},
             tool_calls=msg.get('tool_calls'),
-            stream_state=(ModelStatusCode.SESSION_OUT_OF_LIMIT if finish_reason == 'length' else ModelStatusCode.END),
+            stream_state=choice.get('stream_state')
+            or (ModelStatusCode.SESSION_OUT_OF_LIMIT if finish_reason == 'length' else ModelStatusCode.END),
             finish_reason=finish_reason,
         )