chore: add tracing option to run and group traces under session name (#187)

This commit is contained in:
Zaki Ali
2024-10-24 16:36:06 -07:00
committed by GitHub
parent dccf0ec67c
commit faebda07e3
2 changed files with 14 additions and 5 deletions

View File

@@ -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)

View File

@@ -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]]