Fix agent/trace bugs

Dockerfile

@@ -11,7 +11,7 @@ WORKDIR /app
 COPY . .
 
 RUN poetry install && rm -rf $POETRY_CACHE_DIR
-RUN aliyun-bootstrap -a install
+RUN poetry run aliyun-bootstrap -a install
 
 FROM python:3.11-slim AS prod
 

Dockerfile_gpu

@@ -13,7 +13,7 @@ RUN mv pyproject_gpu.toml pyproject.toml \
     && rm poetry.lock
 
 RUN poetry install && rm -rf $POETRY_CACHE_DIR
-RUN aliyun-bootstrap -a install
+RUN poetry run aliyun-bootstrap -a install
 
 FROM python:3.11-slim AS prod
 

src/pai_rag/app/api/agent_demo.py

@@ -94,9 +94,56 @@
         "arrival_time": "00:30",
         "price": 870,
     },
+    {
+        "train_number": "G85",
+        "from": "上海",
+        "to": "北京",
+        "departure_time": "09:00",
+        "arrival_time": "14:30",
+        "price": 900,
+    },
+    {
+        "train_number": "G87",
+        "from": "上海",
+        "to": "北京",
+        "departure_time": "11:00",
+        "arrival_time": "17:30",
+        "price": 1001,
+    },
+    {
+        "train_number": "G88",
+        "from": "上海",
+        "to": "北京",
+        "departure_time": "13:10",
+        "arrival_time": "17:40",
+        "price": 767,
+    },
+    {
+        "train_number": "G110",
+        "from": "上海",
+        "to": "北京",
+        "departure_time": "18:00",
+        "arrival_time": "23:30",
+        "price": 598,
+    },
 ]
 
 hotels_data = [
+    {
+        "hotel_name": "大地花园酒店",
+        "city": "北京",
+        "price_per_night": 300,
+    },
+    {
+        "hotel_name": "凯悦酒店",
+        "city": "北京",
+        "price_per_night": 1000,
+    },
+    {
+        "hotel_name": "秋果酒店金融街店",
+        "city": "北京",
+        "price_per_night": 500,
+    },
     {
         "hotel_name": "万豪酒店",
         "city": "上海",

src/pai_rag/core/rag_application.py

@@ -53,7 +53,7 @@ async def event_generator_async(
         if token and token != DEFAULT_EMPTY_RESPONSE_GEN:
             chunk = {"delta": token, "is_finished": False}
             content += token
-            yield "data: " + json.dumps(chunk, ensure_ascii=False) + "\n"
+            yield "data: " + json.dumps(chunk, ensure_ascii=False) + "\n\n"
 
     if chat_store:
         chat_store.add_message(
@@ -68,7 +68,7 @@ async def event_generator_async(
 
     yield "data: " + json.dumps(
         last_chunk, default=lambda x: x.dict(), ensure_ascii=False
-    ) + "\n"
+    ) + "\n\n"
 
 
 class RagApplication:

src/pai_rag/integrations/agent/pai/pai_agent.py

@@ -6,7 +6,7 @@
     Optional,
     Type,
 )
-from flask import json
+import json
 from llama_index.agent.openai.step import OpenAIAgentWorker
 from llama_index.core.agent.runner.base import AgentRunner
 from llama_index.core.callbacks import CallbackManager
@@ -26,7 +26,7 @@
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_MAX_FUNCTION_CALLS = 5
+DEFAULT_MAX_FUNCTION_CALLS = 10
 
 
 def get_tools(agent_config: AgentConfig):
@@ -98,7 +98,6 @@ def from_tools(
         max_function_calls: int = DEFAULT_MAX_FUNCTION_CALLS,
         default_tool_choice: str = "auto",
         callback_manager: Optional[CallbackManager] = None,
-        system_prompt: Optional[str] = None,
         prefix_messages: Optional[List[ChatMessage]] = None,
         tool_call_parser: Optional[Callable[[OpenAIToolCall], Dict]] = None,
         **kwargs: Any,
@@ -127,12 +126,14 @@ def from_tools(
                 f"Model name {llm.model} does not support function calling API. "
             )
 
-        if system_prompt is not None:
+        if agent_config.system_prompt is not None:
             if prefix_messages is not None:
                 raise ValueError(
                     "Cannot specify both system_prompt and prefix_messages"
                 )
-            prefix_messages = [ChatMessage(content=system_prompt, role="system")]
+            prefix_messages = [
+                ChatMessage(content=agent_config.system_prompt, role="system")
+            ]
 
         prefix_messages = prefix_messages or []
 

src/pai_rag/integrations/agent/pai/utils/default_tool_description_template.py

@@ -73,3 +73,13 @@
 
     Example uses of this tool include but are not limited to calculating age differences, determining the number of items sold from inventory, working out loan repayments, and any other context where subtraction of numerical values plays a key role.
 """
+
+
+DEFAULT_GET_DATETIME_TOOL = """
+    The get datetime tool is used to retrieve the current date and time. It is particularly useful for tasks that require obtaining the current date and time, such as booking tickets, scheduling appointments, tracking events, or managing time-sensitive tasks.
+    get_current_datetime() -> str
+
+    Return the current date and time as a string using format "%Y-%d-%d %H:%M:%S"..
+
+    This function is essential for tasks that require obtaining the current date and time, such as booking tickets, scheduling appointments, tracking events, or managing time-sensitive tasks.
+"""

src/pai_rag/integrations/agent/pai/utils/tool_utils.py

@@ -1,16 +1,34 @@
+import datetime
 from llama_index.core.tools import FunctionTool
 from pai_rag.integrations.agent.pai.base_tool import ApiTool, PaiAgentDefinition
 from pai_rag.integrations.agent.pai.utils.default_tool_description_template import (
     DEFAULT_CALCULATE_MULTIPLY,
     DEFAULT_CALCULATE_ADD,
     DEFAULT_CALCULATE_DIVIDE,
     DEFAULT_CALCULATE_SUBTRACT,
+    DEFAULT_GET_DATETIME_TOOL,
 )
 import logging
 
 logger = logging.getLogger(__name__)
 
 
+def get_time_tools():
+    def get_current_datetime():
+        """
+        Get the current date and time in the given format.
+        """
+        output_format = "%Y-%m-%d %H:%M:%S"
+        return datetime.datetime.now().strftime(output_format)
+
+    get_current_datetime_tool = FunctionTool.from_defaults(
+        fn=get_current_datetime,
+        name="get_current_datetime",
+        description=DEFAULT_GET_DATETIME_TOOL,
+    )
+    return [get_current_datetime_tool]
+
+
 def get_calculator_tools():
     def multiply(a: int, b: int) -> int:
         """Multiply two integers and returns the result integer"""
@@ -96,6 +114,7 @@ def {function_name}({param_str}):
 
 def get_customized_tools(agent_definition: PaiAgentDefinition):
     tools = []
+    tools.extend(get_time_tools())
 
     # 首先尝试加载python代码
     if agent_definition.python_scripts: