diff --git a/Dockerfile b/Dockerfile index 0a3692fc..ddac4117 100644 --- a/Dockerfile +++ b/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 diff --git a/Dockerfile_gpu b/Dockerfile_gpu index f372a321..e8203f82 100644 --- a/Dockerfile_gpu +++ b/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 diff --git a/src/pai_rag/app/api/agent_demo.py b/src/pai_rag/app/api/agent_demo.py index 4cc57752..e2677718 100644 --- a/src/pai_rag/app/api/agent_demo.py +++ b/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": "上海", diff --git a/src/pai_rag/core/rag_application.py b/src/pai_rag/core/rag_application.py index 8560150e..11839b3f 100644 --- a/src/pai_rag/core/rag_application.py +++ b/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: diff --git a/src/pai_rag/integrations/agent/pai/pai_agent.py b/src/pai_rag/integrations/agent/pai/pai_agent.py index a0234ba6..5edd3ade 100644 --- a/src/pai_rag/integrations/agent/pai/pai_agent.py +++ b/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 is not None and 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 [] diff --git a/src/pai_rag/integrations/agent/pai/utils/default_tool_description_template.py b/src/pai_rag/integrations/agent/pai/utils/default_tool_description_template.py index 2fda5a53..c48219fe 100644 --- a/src/pai_rag/integrations/agent/pai/utils/default_tool_description_template.py +++ b/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. +""" diff --git a/src/pai_rag/integrations/agent/pai/utils/tool_utils.py b/src/pai_rag/integrations/agent/pai/utils/tool_utils.py index aa42020b..1e70d0e5 100644 --- a/src/pai_rag/integrations/agent/pai/utils/tool_utils.py +++ b/src/pai_rag/integrations/agent/pai/utils/tool_utils.py @@ -1,3 +1,4 @@ +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 ( @@ -5,12 +6,29 @@ 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: diff --git a/tests/integrations/agent/test_api_tool.py b/tests/integrations/agent/test_api_tool.py index f7a2e685..e9ad7ef7 100644 --- a/tests/integrations/agent/test_api_tool.py +++ b/tests/integrations/agent/test_api_tool.py @@ -32,6 +32,8 @@ def test_customized_api_tool(): agent_definition = PaiAgentDefinition.model_validate(json_object) tools = get_customized_tools(agent_definition) - assert tools[0].metadata.name == "search_flight_ticket_api" - assert tools[0].metadata.description == "帮助用户获取机票信息，用户需要输入出发地、目的地" - assert tools[0].fn.__name__ == "search_flight_ticket_api" + assert tools[0].metadata.name == "get_current_datetime" + + assert tools[1].metadata.name == "search_flight_ticket_api" + assert tools[1].metadata.description == "帮助用户获取机票信息，用户需要输入出发地、目的地" + assert tools[1].fn.__name__ == "search_flight_ticket_api"