From faebda07e3c5861d4c41149aa4c4507a583c47dc Mon Sep 17 00:00:00 2001
From: Zaki Ali <zaki@squareup.com>
Date: Thu, 24 Oct 2024 16:36:06 -0700
Subject: [PATCH] chore: add tracing option to run and group traces under
 session name (#187)

---
 src/goose/cli/main.py    | 15 +++++++++++----
 src/goose/cli/session.py |  4 +++-
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/goose/cli/main.py b/src/goose/cli/main.py
index a7f484d9..431e26e3 100644
--- a/src/goose/cli/main.py
+++ b/src/goose/cli/main.py
@@ -9,7 +9,7 @@ from ruamel.yaml import YAML
 
 from goose.cli.config import SESSIONS_PATH
 from goose.cli.session import Session
-from goose.toolkit.utils import render_template, parse_plan
+from goose.toolkit.utils import parse_plan, render_template
 from goose.utils import load_plugins
 from goose.utils.autocomplete import SUPPORTED_SHELLS, setup_autocomplete
 from goose.utils.session_file import list_sorted_session_files
@@ -208,7 +208,14 @@ def session_resume(name: Optional[str], profile: str, log_level: str) -> None:
 @click.option("--profile")
 @click.option("--log-level", type=LOG_CHOICE, default="INFO")
 @click.option("--resume-session", is_flag=True, help="Resume the last session if available")
-def run(message_file: Optional[str], profile: str, log_level: str, resume_session: bool = False) -> None:
+@click.option("--tracing", is_flag=True, required=False)
+def run(
+    message_file: Optional[str],
+    profile: str,
+    log_level: str,
+    resume_session: bool = False,
+    tracing: bool = False,
+) -> None:
     """Run a single-pass session with a message from a markdown input file"""
     if message_file:
         with open(message_file, "r") as f:
@@ -220,9 +227,9 @@ def run(message_file: Optional[str], profile: str, log_level: str, resume_sessio
         session_files = get_session_files()
         if session_files:
             name = list(session_files.keys())[0]
-            session = Session(name=name, profile=profile, log_level=log_level)
+            session = Session(name=name, profile=profile, log_level=log_level, tracing=tracing)
     else:
-        session = Session(profile=profile, log_level=log_level)
+        session = Session(profile=profile, log_level=log_level, tracing=tracing)
     session.single_pass(initial_message=initial_message)
 
 
diff --git a/src/goose/cli/session.py b/src/goose/cli/session.py
index 68bdc4ce..9e5d1079 100644
--- a/src/goose/cli/session.py
+++ b/src/goose/cli/session.py
@@ -150,7 +150,6 @@ class Session:
         profile = self.profile_name or "default"
         print(f"[dim]starting session | name:[cyan]{self.name}[/]  profile:[cyan]{profile}[/]")
         print(f"[dim]saving to {self.session_file_path}")
-        print()
 
         # Process initial message
         message = Message.user(initial_message)
@@ -202,6 +201,9 @@ class Session:
     @observe_wrapper()
     def reply(self) -> None:
         """Reply to the last user message, calling tools as needed"""
+        # group all traces under the same session
+        langfuse_context.update_current_trace(session_id=self.name)
+
         # These are the *raw* messages, before the moderator rewrites things
         committed = [self.exchange.messages[-1]]