From 0c3e75445374d5c4a0d76801817873aea0ffc00e Mon Sep 17 00:00:00 2001
From: Hardik Shah <hjshah@fb.com>
Date: Tue, 13 Aug 2024 13:01:05 -0700
Subject: [PATCH] update system prompts to drop new line

---
 llama_toolchain/agentic_system/client.py      | 47 +++++++++++++++++--
 .../meta_reference/system_prompt.py           |  7 ++-
 2 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/llama_toolchain/agentic_system/client.py b/llama_toolchain/agentic_system/client.py
index 71c578e2f..597872e30 100644
--- a/llama_toolchain/agentic_system/client.py
+++ b/llama_toolchain/agentic_system/client.py
@@ -13,7 +13,12 @@ import fire
 
 import httpx
 
-from llama_models.llama3_1.api.datatypes import BuiltinTool, SamplingParams
+from llama_models.llama3_1.api.datatypes import (
+    BuiltinTool,
+    SamplingParams,
+    ToolParamDefinition,
+    UserMessage,
+)
 
 from .api import (
     AgenticSystem,
@@ -87,7 +92,7 @@ class AgenticSystemClient(AgenticSystem):
 
 async def run_main(host: str, port: int):
     # client to test remote impl of agentic system
-    api = await AgenticSystemClient(f"http://{host}:{port}")
+    api = AgenticSystemClient(f"http://{host}:{port}")
 
     tool_definitions = [
         AgenticSystemToolDefinition(
@@ -103,6 +108,19 @@ async def run_main(host: str, port: int):
             tool_name=BuiltinTool.code_interpreter,
         ),
     ]
+    tool_definitions += [
+        AgenticSystemToolDefinition(
+            tool_name="custom_tool",
+            description="a custom tool",
+            parameters={
+                "param1": ToolParamDefinition(
+                    param_type="str",
+                    description="a string parameter",
+                    required=True,
+                )
+            },
+        )
+    ]
 
     create_request = AgenticSystemCreateRequest(
         model="Meta-Llama3.1-8B-Instruct",
@@ -118,8 +136,29 @@ async def run_main(host: str, port: int):
     )
 
     create_response = await api.create_agentic_system(create_request)
-    print(create_response)
-    # TODO: Add chat session / turn apis to test e2e
+    print("Create Response -->", create_response)
+
+    session_response = await api.create_agentic_system_session(
+        AgenticSystemSessionCreateRequest(
+            system_id=create_response.system_id,
+            session_name="test_session",
+        )
+    )
+    print("Session Response -->", session_response)
+
+    turn_response = api.create_agentic_system_turn(
+        AgenticSystemTurnCreateRequest(
+            system_id=create_response.system_id,
+            session_id=session_response.session_id,
+            messages=[
+                UserMessage(content="Who are you?"),
+            ],
+            stream=False,
+        )
+    )
+    print("Turn Response -->")
+    async for chunk in turn_response:
+        print(chunk)
 
 
 def main(host: str, port: int):
diff --git a/llama_toolchain/agentic_system/meta_reference/system_prompt.py b/llama_toolchain/agentic_system/meta_reference/system_prompt.py
index c8c616285..d51e53a82 100644
--- a/llama_toolchain/agentic_system/meta_reference/system_prompt.py
+++ b/llama_toolchain/agentic_system/meta_reference/system_prompt.py
@@ -34,13 +34,13 @@ def get_agentic_prefix_messages(
             ]
         )
         if tool_str:
-            content += f"Tools: {tool_str}\n"
+            content += f"Tools: {tool_str}"
 
     current_date = datetime.now()
     formatted_date = current_date.strftime("%d %B %Y")
     date_str = f"""
 Cutting Knowledge Date: December 2023
-Today Date: {formatted_date}\n\n"""
+Today Date: {formatted_date}\n"""
     content += date_str
 
     if custom_tools:
@@ -49,7 +49,7 @@ Today Date: {formatted_date}\n\n"""
 
     # TODO: Replace this hard coded message with instructions coming in the request
     if False:
-        content += "You are a helpful Assistant."
+        content += "\nYou are a helpful Assistant."
 
     messages.append(SystemMessage(content=content))
     return messages
@@ -76,7 +76,6 @@ Reminder:
 - Required parameters MUST be specified
 - Only call one function at a time
 - Put the entire function call reply on one line
-
 """
     return content