snowglobe 0.4.11__tar.gz → 0.4.13__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (26) hide show
  1. {snowglobe-0.4.11/src/snowglobe.egg-info → snowglobe-0.4.13}/PKG-INFO +3 -2
  2. {snowglobe-0.4.11 → snowglobe-0.4.13}/README.md +1 -1
  3. {snowglobe-0.4.11 → snowglobe-0.4.13}/pyproject.toml +2 -1
  4. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/app.py +78 -11
  5. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/cli.py +298 -76
  6. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/cli_utils.py +40 -13
  7. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/models.py +18 -2
  8. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/telemetry.py +7 -3
  9. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/utils.py +7 -4
  10. {snowglobe-0.4.11 → snowglobe-0.4.13/src/snowglobe.egg-info}/PKG-INFO +3 -2
  11. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe.egg-info/requires.txt +1 -0
  12. {snowglobe-0.4.11 → snowglobe-0.4.13}/tests/test_cli.py +23 -20
  13. {snowglobe-0.4.11 → snowglobe-0.4.13}/LICENSE +0 -0
  14. {snowglobe-0.4.11 → snowglobe-0.4.13}/setup.cfg +0 -0
  15. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/__init__.py +0 -0
  16. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/config.py +0 -0
  17. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/project_manager.py +0 -0
  18. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe/client/src/stats.py +0 -0
  19. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe.egg-info/SOURCES.txt +0 -0
  20. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe.egg-info/dependency_links.txt +0 -0
  21. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe.egg-info/entry_points.txt +0 -0
  22. {snowglobe-0.4.11 → snowglobe-0.4.13}/src/snowglobe.egg-info/top_level.txt +0 -0
  23. {snowglobe-0.4.11 → snowglobe-0.4.13}/tests/test_app.py +0 -0
  24. {snowglobe-0.4.11 → snowglobe-0.4.13}/tests/test_config.py +0 -0
  25. {snowglobe-0.4.11 → snowglobe-0.4.13}/tests/test_heartbeat.py +0 -0
  26. {snowglobe-0.4.11 → snowglobe-0.4.13}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: snowglobe
3
- Version: 0.4.11
3
+ Version: 0.4.13
4
4
  Summary: client server for usage with snowglobe experiments
5
5
  Author-email: Guardrails AI <contact@guardrailsai.com>
6
6
  License: MIT License
@@ -31,6 +31,7 @@ License-File: LICENSE
31
31
  Requires-Dist: typer>=0.9.0
32
32
  Requires-Dist: fastapi>=0.115.6
33
33
  Requires-Dist: uvicorn[standard]>=0.23.0
34
+ Requires-Dist: openai>=1.55.2
34
35
  Requires-Dist: requests>=2.31.0
35
36
  Requires-Dist: pydantic>=2.11.5
36
37
  Requires-Dist: APScheduler==4.0.0a6
@@ -87,7 +88,7 @@ import os
87
88
 
88
89
  client = OpenAI(api_key=os.getenv("SNOWGLOBE_API_KEY"))
89
90
 
90
- def process_scenario(request: CompletionRequest) -> CompletionFunctionOutputs:
91
+ def completion_fn(request: CompletionRequest) -> CompletionFunctionOutputs:
91
92
  """
92
93
  Process a scenario request from Snowglobe.
93
94
 
@@ -42,7 +42,7 @@ import os
42
42
 
43
43
  client = OpenAI(api_key=os.getenv("SNOWGLOBE_API_KEY"))
44
44
 
45
- def process_scenario(request: CompletionRequest) -> CompletionFunctionOutputs:
45
+ def completion_fn(request: CompletionRequest) -> CompletionFunctionOutputs:
46
46
  """
47
47
  Process a scenario request from Snowglobe.
48
48
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "snowglobe"
3
- version = "0.4.11"
3
+ version = "0.4.13"
4
4
  authors = [
5
5
  {name = "Guardrails AI", email = "contact@guardrailsai.com"}
6
6
  ]
@@ -14,6 +14,7 @@ dependencies = [
14
14
  "typer>=0.9.0",
15
15
  "fastapi>=0.115.6",
16
16
  "uvicorn[standard]>=0.23.0",
17
+ "openai>=1.55.2",
17
18
  "requests>=2.31.0",
18
19
  "pydantic>=2.11.5",
19
20
  "APScheduler==4.0.0a6",
@@ -26,7 +26,13 @@ from snowglobe.client.src.telemetry import trace_completion_fn, trace_risk_evalu
26
26
 
27
27
  from .cli_utils import info, shutdown_manager
28
28
  from .config import config, get_api_key_or_raise
29
- from .models import CompletionFunctionOutputs, CompletionRequest, RiskEvaluationRequest, SnowglobeData, SnowglobeMessage
29
+ from .models import (
30
+ CompletionFunctionOutputs,
31
+ CompletionRequest,
32
+ RiskEvaluationRequest,
33
+ SnowglobeData,
34
+ SnowglobeMessage,
35
+ )
30
36
  from .stats import initialize_stats, track_batch_completion
31
37
  from .utils import fetch_experiments, fetch_messages
32
38
 
@@ -130,13 +136,16 @@ async def process_application_heartbeat(app_id):
130
136
  try:
131
137
  prompt = "Hello from Snowglobe!"
132
138
  test_request = CompletionRequest(
133
- messages=[SnowglobeMessage(
139
+ messages=[
140
+ SnowglobeMessage(
134
141
  role="user",
135
142
  content=prompt,
136
143
  snowglobe_data=SnowglobeData(
137
144
  conversation_id="test", test_id="test"
138
145
  ),
139
- )])
146
+ )
147
+ ]
148
+ )
140
149
  heartbeat_id = uuid.uuid4().hex
141
150
  agent = apps.get(app_id, {})
142
151
  agent_name = agent.get("name", "")
@@ -157,7 +166,11 @@ async def process_application_heartbeat(app_id):
157
166
  )
158
167
  async def run_completion_fn(completion_request: CompletionRequest):
159
168
  if asyncio.iscoroutinefunction(completion_fn):
160
- response = await completion_fn(completion_request)
169
+ try:
170
+ asyncio.get_running_loop()
171
+ response = await completion_fn(completion_request)
172
+ except RuntimeError:
173
+ response = asyncio.run(completion_fn(completion_request))
161
174
  else:
162
175
  response = completion_fn(completion_request)
163
176
  return response
@@ -250,7 +263,11 @@ async def process_risk_evaluation(test, risk_name, simulation_name, agent_name):
250
263
  )
251
264
  async def run_risk_evaluation_fn(risk_evaluation_request: RiskEvaluationRequest):
252
265
  if asyncio.iscoroutinefunction(risks[risk_name]):
253
- risk_evaluation = await risks[risk_name](risk_evaluation_request)
266
+ try:
267
+ asyncio.get_running_loop()
268
+ risk_evaluation = await risks[risk_name](risk_evaluation_request)
269
+ except RuntimeError:
270
+ risk_evaluation = asyncio.run(risks[risk_name](risk_evaluation_request))
254
271
  else:
255
272
  risk_evaluation = risks[risk_name](risk_evaluation_request)
256
273
  return risk_evaluation
@@ -312,7 +329,11 @@ async def process_test(test, completion_fn, app_id, simulation_name):
312
329
  )
313
330
  async def run_completion_fn(completion_request: CompletionRequest):
314
331
  if asyncio.iscoroutinefunction(completion_fn):
315
- completionOutput = await completion_fn(completion_request)
332
+ try:
333
+ asyncio.get_running_loop()
334
+ completionOutput = await completion_fn(completion_request)
335
+ except RuntimeError:
336
+ completionOutput = asyncio.run(completion_fn(completion_request))
316
337
  else:
317
338
  completionOutput = completion_fn(completion_request)
318
339
  return completionOutput
@@ -791,14 +812,30 @@ async def lifespan(app: FastAPI):
791
812
  finally:
792
813
  sys.path = sys_path_backup
793
814
 
794
- if not hasattr(agent_module, "process_scenario"):
815
+ completion_fn = None
816
+ if hasattr(agent_module, "completion"):
817
+ completion_fn = agent_module.completion
818
+ elif hasattr(agent_module, "acompletion"):
819
+ completion_fn = agent_module.acompletion
820
+ # Check for legacy function names
821
+ elif hasattr(agent_module, "completion_fn"):
822
+ completion_fn = agent_module.completion_fn
823
+ LOGGER.warning(
824
+ f"Agent {filename} uses deprecated function 'completion_fn'. Please rename to 'completion'"
825
+ )
826
+ elif hasattr(agent_module, "process_scenario"):
827
+ completion_fn = agent_module.process_scenario
828
+ LOGGER.warning(
829
+ f"Agent {filename} uses deprecated function 'process_scenario'. Please rename to 'completion' or 'acompletion'"
830
+ )
831
+ else:
795
832
  LOGGER.warning(
796
- f"Agent {filename} does not have a process_scenario function"
833
+ f"Agent {filename} does not have a completion or acompletion function"
797
834
  )
798
835
  continue
799
836
 
800
837
  apps[app_id] = {
801
- "completion_fn": agent_module.process_scenario,
838
+ "completion_fn": completion_fn,
802
839
  "name": app_name,
803
840
  }
804
841
 
@@ -836,17 +873,47 @@ async def lifespan(app: FastAPI):
836
873
  )
837
874
  sg_connect = importlib.util.module_from_spec(spec)
838
875
  spec.loader.exec_module(sg_connect)
839
- if hasattr(sg_connect, "completion_fn"):
876
+ if hasattr(sg_connect, "completion"):
877
+ apps[config.APPLICATION_ID] = {
878
+ "completion_fn": sg_connect.completion,
879
+ "name": "Legacy Single Application",
880
+ }
881
+ LOGGER.info(
882
+ f"Loaded legacy application with ID {config.APPLICATION_ID} for completions."
883
+ )
884
+ elif hasattr(sg_connect, "acompletion"):
885
+ apps[config.APPLICATION_ID] = {
886
+ "completion_fn": sg_connect.acompletion,
887
+ "name": "Legacy Single Application",
888
+ }
889
+ LOGGER.info(
890
+ f"Loaded legacy application with ID {config.APPLICATION_ID} for completions."
891
+ )
892
+ elif hasattr(sg_connect, "completion_fn"):
840
893
  apps[config.APPLICATION_ID] = {
841
894
  "completion_fn": sg_connect.completion_fn,
842
895
  "name": "Legacy Single Application",
843
896
  }
897
+ LOGGER.warning(
898
+ f"Legacy application with ID {config.APPLICATION_ID} uses deprecated function 'completion_fn'. Please rename to 'completion'"
899
+ )
900
+ LOGGER.info(
901
+ f"Loaded legacy application with ID {config.APPLICATION_ID} for completions."
902
+ )
903
+ elif hasattr(sg_connect, "process_scenario"):
904
+ apps[config.APPLICATION_ID] = {
905
+ "completion_fn": sg_connect.process_scenario,
906
+ "name": "Legacy Single Application",
907
+ }
908
+ LOGGER.warning(
909
+ f"Legacy application with ID {config.APPLICATION_ID} uses deprecated function 'process_scenario'. Please rename to 'completion' or 'acompletion'"
910
+ )
844
911
  LOGGER.info(
845
912
  f"Loaded legacy application with ID {config.APPLICATION_ID} for completions."
846
913
  )
847
914
  else:
848
915
  LOGGER.error(
849
- f"Legacy application with ID {config.APPLICATION_ID} does not have a completion_fn."
916
+ f"Legacy application with ID {config.APPLICATION_ID} does not have a completion or acompletion function."
850
917
  )
851
918
  except Exception as e:
852
919
  LOGGER.error(f"Error loading applications: {e}")
@@ -31,7 +31,7 @@ from .cli_utils import (
31
31
  graceful_shutdown,
32
32
  info,
33
33
  select_application_interactive,
34
- select_stateful_interactive,
34
+ select_template_interactive,
35
35
  spinner,
36
36
  success,
37
37
  warning,
@@ -196,7 +196,7 @@ def test(
196
196
  if "default template response" in conn_message.lower():
197
197
  info("This is expected with the default template.")
198
198
  info(
199
- "Please implement your application logic in the process_scenario function."
199
+ "Please implement your application logic in the completion or acompletion function."
200
200
  )
201
201
  else:
202
202
  info("Check your implementation and try again.")
@@ -214,10 +214,6 @@ def init(
214
214
  "-f",
215
215
  help="Path or filename (within project) for the agent wrapper",
216
216
  ),
217
- # option for stateful agent
218
- stateful: bool = typer.Option(
219
- False, "--stateful", help="Initialize a stateful agent template"
220
- ),
221
217
  ):
222
218
  """
223
219
  Initialize a new Snowglobe agent in the current directory
@@ -269,8 +265,8 @@ def init(
269
265
 
270
266
  success(f"Selected application: {app_name}")
271
267
 
272
- # prompt user to confirm if stateful agent
273
- user_stateful = select_stateful_interactive(stateful)
268
+ # prompt user to select template type
269
+ template_type = select_template_interactive()
274
270
  # Set up project structure
275
271
  with spinner("Setting up project structure"):
276
272
  pm.ensure_project_structure()
@@ -310,11 +306,16 @@ def init(
310
306
  file_path = pm.project_root / filename
311
307
 
312
308
  success(f"Using filename: {filename}")
313
- # use the appopriate template based on stateful option
314
- if user_stateful:
315
- snowglobe_connect_template = stateful_snowglobe_connect_template
309
+ # use the appropriate template based on template type
310
+ if template_type == "sync":
311
+ snowglobe_connect_template = sync_snowglobe_connect_template
312
+ elif template_type == "async":
313
+ snowglobe_connect_template = async_snowglobe_connect_template
314
+ elif template_type == "socket":
315
+ snowglobe_connect_template = socket_snowglobe_connect_template
316
316
  else:
317
- snowglobe_connect_template = stateless_snowglobe_connect_template
317
+ # Default fallback
318
+ snowglobe_connect_template = sync_snowglobe_connect_template
318
319
 
319
320
  # Ensure parent directories exist if a subpath was provided
320
321
  os.makedirs(file_path.parent, exist_ok=True)
@@ -334,13 +335,17 @@ def init(
334
335
  console.print(f"[dim] {filename}\t- Your agent wrapper[/dim]")
335
336
 
336
337
  console.print()
338
+ console.print("📁 Your connection template is available at:")
339
+ console.print(f"[bold blue]{'*' * (len(filename) + 4)}[/bold blue]")
340
+ console.print(f"[bold blue]* {filename} *[/bold blue]")
341
+ console.print(f"[bold blue]{'*' * (len(filename) + 4)}[/bold blue]")
342
+ console.print(
343
+ "[bold yellow]Please change the code in the completion or acompletion function to implement your application logic.[/bold yellow]\n"
344
+ )
337
345
  info("Next steps:")
338
- console.print("1. Edit the process_scenario function in your agent file:")
339
- console.print(f" [bold cyan]{filename}[/bold cyan]")
340
- console.print("2. Implement your application logic")
341
- console.print("3. Test your agent:")
346
+ console.print("1. Test your agent:")
342
347
  console.print(" [bold green]snowglobe-connect test[/bold green]")
343
- console.print("4. Start the client:")
348
+ console.print("2. Start the client:")
344
349
  console.print(" [bold green]snowglobe-connect start[/bold green]")
345
350
 
346
351
  console.print()
@@ -368,12 +373,29 @@ def test_agent_wrapper(filename: str, app_id: str, app_name: str) -> Tuple[bool,
368
373
  finally:
369
374
  sys.path = sys_path_backup
370
375
 
371
- if not hasattr(agent_module, "process_scenario"):
372
- return False, "process_scenario function not found"
376
+ completion_fn = None
377
+
378
+ # Check for preferred function names first
379
+ if hasattr(agent_module, "completion"):
380
+ completion_fn = agent_module.completion
381
+ elif hasattr(agent_module, "acompletion"):
382
+ completion_fn = agent_module.acompletion
383
+ # Check for legacy function names
384
+ elif hasattr(agent_module, "completion_fn"):
385
+ completion_fn = agent_module.completion_fn
386
+ warning(
387
+ "Function 'completion_fn' is deprecated. Please rename to 'completion'"
388
+ )
389
+ elif hasattr(agent_module, "process_scenario"):
390
+ completion_fn = agent_module.process_scenario
391
+ warning(
392
+ "Function 'process_scenario' is deprecated. Please rename to 'completion' or 'acompletion'"
393
+ )
394
+ else:
395
+ return False, "completion or acompletion function not found"
373
396
 
374
- process_scenario = agent_module.process_scenario
375
- if not callable(process_scenario):
376
- return False, "process_scenario is not callable"
397
+ if not callable(completion_fn):
398
+ return False, "completion function is not callable"
377
399
 
378
400
  # Test with a simple request
379
401
 
@@ -399,14 +421,18 @@ def test_agent_wrapper(filename: str, app_id: str, app_name: str) -> Tuple[bool,
399
421
  simulation_name=f"{app_name} CLI Test",
400
422
  span_type="snowglobe/cli-test",
401
423
  )
402
- async def run_process_scenario(completion_request: CompletionRequest):
403
- if asyncio.iscoroutinefunction(process_scenario):
404
- response = asyncio.run(process_scenario(completion_request))
424
+ async def run_completion_fn(completion_request: CompletionRequest):
425
+ if asyncio.iscoroutinefunction(completion_fn):
426
+ try:
427
+ asyncio.get_running_loop()
428
+ response = await completion_fn(completion_request)
429
+ except RuntimeError:
430
+ response = asyncio.run(completion_fn(completion_request))
405
431
  else:
406
- response = process_scenario(completion_request)
432
+ response = completion_fn(completion_request)
407
433
  return response
408
434
 
409
- response = asyncio.run(run_process_scenario(test_request))
435
+ response = asyncio.run(run_completion_fn(test_request))
410
436
 
411
437
  if hasattr(response, "response") and isinstance(response.response, str):
412
438
  if response.response == "Your response here":
@@ -473,84 +499,127 @@ def enhanced_error_handler(status_code: int, operation: str = "operation") -> No
473
499
  info("Please try again or contact support if the issue persists")
474
500
 
475
501
 
476
- stateless_snowglobe_connect_template = """from snowglobe.client import CompletionRequest, CompletionFunctionOutputs
477
-
502
+ sync_snowglobe_connect_template = """from snowglobe.client import CompletionRequest, CompletionFunctionOutputs
503
+ from openai import OpenAI
504
+ import os
505
+ client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
478
506
 
479
- def process_scenario(request: CompletionRequest) -> CompletionFunctionOutputs:
507
+ def completion(request: CompletionRequest) -> CompletionFunctionOutputs:
480
508
  \"\"\"
481
509
  Process a scenario request from Snowglobe.
482
510
 
483
- This function is called by the Snowglobe client to process requests. It should return a
511
+ This function is called by the Snowglobe client to process test requests. It should return a
484
512
  CompletionFunctionOutputs object with the response content.
513
+
514
+ Args:
515
+ request (CompletionRequest): The request object containing messages for the test.
485
516
 
486
- Example CompletionRequest:
487
- CompletionRequest(
488
- messages=[
489
- SnowglobeMessage(role="user", content="Hello, how are you?", snowglobe_data=None),
490
- ]
517
+ Returns:
518
+ CompletionFunctionOutputs: The response object with the generated content.
519
+ \"\"\"
520
+
521
+ # Process the request using the messages. Example using OpenAI:
522
+ messages = request.to_openai_messages(system_prompt="You are a helpful assistant.")
523
+ response = client.chat.completions.create(
524
+ model="gpt-4o-mini",
525
+ messages=messages
491
526
  )
527
+ return CompletionFunctionOutputs(response=response.choices[0].message.content)
528
+ """
492
529
 
493
- Example CompletionFunctionOutputs:
494
- CompletionFunctionOutputs(response="This is a string response from your application")
530
+ async_snowglobe_connect_template = """from snowglobe.client import CompletionRequest, CompletionFunctionOutputs
531
+ from openai import AsyncOpenAI
532
+ import os
533
+ client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
534
+
535
+ async def acompletion(request: CompletionRequest) -> CompletionFunctionOutputs:
536
+ \"\"\"
537
+ Process a scenario request from Snowglobe.
538
+
539
+ This function is called by the Snowglobe client to process test requests. It should return a
540
+ CompletionFunctionOutputs object with the response content.
495
541
 
496
542
  Args:
497
- request (CompletionRequest): The request object containing the messages.
543
+ request (CompletionRequest): The request object containing messages for the test.
498
544
 
499
545
  Returns:
500
546
  CompletionFunctionOutputs: The response object with the generated content.
501
547
  \"\"\"
502
548
 
503
- # Process the request using the messages. Example:
504
- # messages = request.to_openai_messages()
505
- # response = client.chat.completions.create(
506
- # model="gpt-4o-mini",
507
- # messages=messages
508
- # )
509
- return CompletionFunctionOutputs(response="Your response here")
510
-
549
+ # Process the request using the messages. Example using OpenAI:
550
+ messages = request.to_openai_messages(system_prompt="You are a helpful assistant.")
551
+ response = await client.chat.completions.create(
552
+ model="gpt-4o-mini",
553
+ messages=messages
554
+ )
555
+ return CompletionFunctionOutputs(response=response.choices[0].message.content)
511
556
  """
512
557
 
513
- stateful_snowglobe_connect_template = """
514
- # This file is auto-generated by the API client.
515
-
558
+ socket_snowglobe_connect_template = """
516
559
  from snowglobe.client import CompletionRequest, CompletionFunctionOutputs
517
560
  import logging
518
561
  import websockets
519
562
  import json
563
+ from openai import AsyncOpenAI
564
+
520
565
  LOGGER = logging.getLogger(__name__)
521
566
  socket_cache = {}
522
- async def completion_fn(request: CompletionRequest) -> CompletionFunctionOutputs:
523
- # for debugging purposes
524
- # print(f"Received request: {request.messages}")
525
-
526
- completion_messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
567
+ openai_client = AsyncOpenAI()
568
+
569
+ async def acompletion(request: CompletionRequest) -> CompletionFunctionOutputs:
570
+ \"\"\"
571
+ When dealing with a realtime socket, we need to create a socket for each conversation.
572
+ We store the socket in a cache and reuse it for the same conversation_id so that we can maintain the conversation context.
573
+ Swap out the websocket client for your preferred realtime client.
574
+
575
+ Args:
576
+ request (CompletionRequest): The request object containing messages for the test.
577
+
578
+ Returns:
579
+ CompletionFunctionOutputs: The response object with the generated content.
580
+ \"\"\"
581
+ conversation_id = request.get_conversation_id()
527
582
 
528
- # check the socket cache for a socket for this conversation_id
529
- conversation_id = request.messages[0].snowglobe_data.conversation_id
530
- if conversation_id in socket_cache:
531
- socket = socket_cache[conversation_id]
532
- else:
533
- # create a new socket connection
534
- # this is talking to a local socket/stateful server that you implemented
535
- # update this or implement a way to connect to your actual server
536
- socket = await websockets.connect("ws://localhost:9000/ws")
583
+ if conversation_id not in socket_cache:
584
+ socket = await websockets.connect(
585
+ "wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01&modalities=text",
586
+ additional_headers={
587
+ "Authorization": f"Bearer {openai_client.api_key}",
588
+ "OpenAI-Beta": "realtime=v1"
589
+ }
590
+ )
537
591
  socket_cache[conversation_id] = socket
592
+ else:
593
+ socket = socket_cache[conversation_id]
594
+
595
+ # Send user message
596
+ messages = request.to_openai_messages()
597
+ user_message = messages[-1]["content"]
538
598
 
539
- # send the request to the socket
540
599
  await socket.send(json.dumps({
541
- "messages": completion_messages,
542
- "conversation_id": conversation_id
600
+ "type": "conversation.item.create",
601
+ "session": {
602
+ "modalities": ["text"], # Only text, no audio
603
+ },
604
+ "item": {
605
+ "type": "message",
606
+ "role": "user",
607
+ "content": [{"type": "input_text", "text": user_message}]
608
+ }
543
609
  }))
544
-
545
- # wait for the response from the socket
546
- response = await socket.recv()
547
- response_data = json.loads(response)
548
-
549
- # for debugging purposes
550
- # print(f"Received response from socket: {response_data}")
551
610
 
552
- # Example implementation, replace with actual logic
553
- return CompletionFunctionOutputs(response=response_data.messages[0]["content"])
611
+ await socket.send(json.dumps({"type": "response.create"}))
612
+
613
+ # Get response
614
+ response_content = ""
615
+ async for message in socket:
616
+ data = json.loads(message)
617
+ if data.get("type") == "response.audio_transcript.delta":
618
+ response_content += data.get("delta", "")
619
+ elif data.get("type") == "response.done":
620
+ break
621
+
622
+ return CompletionFunctionOutputs(response=response_content)
554
623
  """
555
624
 
556
625
 
@@ -741,6 +810,159 @@ def auth(
741
810
  _poll_for_api_key(rc_path)
742
811
 
743
812
 
813
+ @cli_app.command()
814
+ def manage():
815
+ """
816
+ Manage saved app connections interactively
817
+ """
818
+ print("\n🔗 Manage Chatbot Connections\n")
819
+
820
+ pm = get_project_manager()
821
+
822
+ while True:
823
+ # Load current connections
824
+ agents = pm.list_agents()
825
+
826
+ if not agents:
827
+ print("💡 No connections found")
828
+ print("Run 'snowglobe-connect init' to create connections")
829
+ return
830
+
831
+ # Display connections
832
+ print(f"📱 Chatbot Connections ({len(agents)} total)")
833
+ print("-" * 60)
834
+ print(f"{'#':<3} {'File':<20} {'Chatbot Name':<20} {'Created':<12}")
835
+ print("-" * 60)
836
+
837
+ for i, (filename, agent_info) in enumerate(agents, 1):
838
+ app_name = agent_info.get("name", "Unknown")
839
+ created = agent_info.get("created", "Unknown")
840
+ # Format the created date
841
+ if "T" in created:
842
+ created = created.split("T")[0]
843
+
844
+ print(f"{i:<3} {filename:<20} {app_name:<20} {created:<12}")
845
+
846
+ print("-" * 60)
847
+ print("Commands:")
848
+ console.print(
849
+ f"- [bold green]1-{len(agents)}[/bold green] View connection details"
850
+ )
851
+ console.print("- [bold green]d[/bold green] Delete a connection")
852
+ console.print("- [bold green]q[/bold green] Quit")
853
+
854
+ try:
855
+ choice = input("").strip().lower()
856
+
857
+ if choice == "q":
858
+ print("✅ Session ended")
859
+ return
860
+ elif choice == "d":
861
+ # Delete mode
862
+ if not _handle_delete_mode(pm, agents):
863
+ continue
864
+ elif choice.isdigit():
865
+ idx = int(choice)
866
+ if 1 <= idx <= len(agents):
867
+ filename, agent_info = agents[idx - 1]
868
+ _show_connection_details(filename, agent_info)
869
+ else:
870
+ print(f"❌ Please choose between 1 and {len(agents)}")
871
+ else:
872
+ print("❌ Invalid choice. Use a number, 'd' to delete, or 'q' to quit")
873
+
874
+ except (KeyboardInterrupt, EOFError):
875
+ print("\n✅ Session ended")
876
+ return
877
+
878
+
879
+ def _handle_delete_mode(pm, agents):
880
+ """Handle connection deletion workflow"""
881
+ print("\n🗑️ Delete Connection")
882
+ print("Select a connection to delete:")
883
+
884
+ # Show numbered list for deletion
885
+ for i, (filename, agent_info) in enumerate(agents, 1):
886
+ app_name = agent_info.get("name", "Unknown")
887
+ print(f" {i}. {filename} ({app_name})")
888
+
889
+ try:
890
+ choice = (
891
+ input("\nSelect connection to delete (number or 'c' to cancel): ")
892
+ .strip()
893
+ .lower()
894
+ )
895
+
896
+ if choice == "c":
897
+ return True # Continue with main menu
898
+ elif choice.isdigit():
899
+ idx = int(choice)
900
+ if 1 <= idx <= len(agents):
901
+ filename, agent_info = agents[idx - 1]
902
+ app_name = agent_info.get("name", "Unknown")
903
+
904
+ # Confirmation
905
+ print("\n⚠️ You are about to delete:")
906
+ print(f" File: {filename}")
907
+ print(f" Chatbot: {app_name}")
908
+
909
+ confirm = (
910
+ input("\nAre you sure you want to delete this connection? (y/N): ")
911
+ .strip()
912
+ .lower()
913
+ )
914
+
915
+ if confirm == "y":
916
+ # Remove from mapping
917
+ pm.remove_agent_mapping(filename)
918
+
919
+ # Ask if they want to delete the file too
920
+ file_path = pm.project_root / filename
921
+ if file_path.exists():
922
+ delete_file = (
923
+ input(
924
+ f"Also delete the file '{filename}' from disk? (y/N): "
925
+ )
926
+ .strip()
927
+ .lower()
928
+ )
929
+ if delete_file == "y":
930
+ try:
931
+ file_path.unlink()
932
+ print(f"✅ Deleted connection and file: {filename}")
933
+ except Exception as e:
934
+ print(f"❌ Failed to delete file: {e}")
935
+ print(f"✅ Connection mapping removed: {filename}")
936
+ else:
937
+ print(f"✅ Connection mapping removed: {filename}")
938
+ else:
939
+ print(f"✅ Connection mapping removed: {filename}")
940
+ else:
941
+ print("💡 Delete cancelled")
942
+
943
+ return True # Continue with main menu
944
+ else:
945
+ print(f"❌ Please choose between 1 and {len(agents)}")
946
+ return False # Stay in delete mode
947
+ else:
948
+ print("❌ Invalid choice")
949
+ return False # Stay in delete mode
950
+
951
+ except (KeyboardInterrupt, EOFError):
952
+ print("\n💡 Delete cancelled")
953
+ return True # Continue with main menu
954
+
955
+
956
+ def _show_connection_details(filename, agent_info):
957
+ """Show detailed information about a connection"""
958
+ print("\n📋 Connection Details")
959
+ print(f" File: {filename}")
960
+ print(f" Chatbot Name: {agent_info.get('name', 'Unknown')}")
961
+ print(f" UUID: {agent_info.get('uuid', 'Unknown')}")
962
+ print(f" Created: {agent_info.get('created', 'Unknown')}")
963
+ input("\nPress Enter to continue...")
964
+
965
+
744
966
  @cli_app.command()
745
967
  def start(
746
968
  verbose: bool = typer.Option(
@@ -311,22 +311,49 @@ def check_auth_status() -> Tuple[bool, str, Dict[str, Any]]:
311
311
  return False, f"Connection error: {str(e)}", {}
312
312
 
313
313
 
314
- def select_stateful_interactive(
315
- stateful: bool = False,
316
- ) -> bool:
317
- """Interactive prompt to confirm if the agent is stateful"""
314
+ def select_template_interactive(
315
+ template: Optional[str] = None,
316
+ ) -> str:
317
+ """Interactive prompt to select template type"""
318
318
  if cli_state.json_output:
319
- # For JSON mode, just return the default stateful value
320
- return stateful
321
- info(
322
- "Some stateful agents such as ones that maintain communication over a websocket or convo specific completion endpoint require stateful integration."
319
+ # For JSON mode, return default or provided template
320
+ return template or "sync"
321
+
322
+ console.print("\n[bold cyan]📋 Select Integration Template:[/bold cyan]")
323
+ console.print(
324
+ "1. [bold green]Sync[/bold green] - Simple synchronous completion function"
325
+ )
326
+ console.print(
327
+ " 💡 Standard API calls, simple request-response patterns. Ex: OpenAI Client"
328
+ )
329
+ console.print(
330
+ "2. [bold yellow]Async[/bold yellow] - Asynchronous completion function"
323
331
  )
324
- info(
325
- "If your agent takes messages and completions on a single completion endpoint regardless of context, you can answer no to the following question."
332
+ console.print(
333
+ " 💡 Non-blocking operations, concurrent processing. Ex: OpenAI AsyncClient"
326
334
  )
327
- if Confirm.ask("Would you like to initialize a stateful application?"):
328
- return True
329
- return False
335
+ console.print(
336
+ "3. [bold magenta]Socket[/bold magenta] - Real-time WebSocket/stateful connection"
337
+ )
338
+ console.print(
339
+ " 💡 Conversational agents, real-time streaming, stateful interactions. Ex: OpenAI Realtime Client"
340
+ )
341
+
342
+ from rich.prompt import Prompt
343
+
344
+ while True:
345
+ choice = Prompt.ask(
346
+ "\n[bold]Choose template[/bold]", choices=["1", "2", "3"], default="1"
347
+ )
348
+
349
+ if choice == "1":
350
+ return "sync"
351
+ elif choice == "2":
352
+ return "async"
353
+ elif choice == "3":
354
+ return "socket"
355
+ else:
356
+ error("Please choose 1, 2, or 3")
330
357
 
331
358
 
332
359
  def select_application_interactive(
@@ -21,9 +21,25 @@ class CompletionFunctionOutputs(BaseModel):
21
21
  class CompletionRequest(BaseModel):
22
22
  messages: List[SnowglobeMessage]
23
23
 
24
- def to_openai_messages(self) -> List[Dict]:
24
+ def to_openai_messages(self, system_prompt: Optional[str] = None) -> List[Dict]:
25
25
  """Return a list of OpenAI messages from the Snowglobe messages"""
26
- return [{"role": msg.role, "content": msg.content} for msg in self.messages]
26
+ oai_messages = []
27
+
28
+ if system_prompt:
29
+ oai_messages.append({"role": "system", "content": system_prompt})
30
+
31
+ oai_messages.extend(
32
+ [{"role": msg.role, "content": msg.content} for msg in self.messages]
33
+ )
34
+ return oai_messages
35
+
36
+ def get_prompt(self) -> str:
37
+ """Return the prompt from the Snowglobe messages"""
38
+ return self.to_openai_messages(system_prompt=None)[-1]["content"]
39
+
40
+ def get_conversation_id(self) -> str:
41
+ """Return the conversation id from the Snowglobe messages"""
42
+ return self.messages[0].snowglobe_data.conversation_id
27
43
 
28
44
 
29
45
  class RiskEvaluationRequest(BaseModel):
@@ -33,7 +33,9 @@ def trace_completion_fn(
33
33
  current_user = w.current_user.me()
34
34
 
35
35
  formatted_sim_name = simulation_name.lower().replace(" ", "_")
36
- default_experiment_name = f"/Users/{current_user.user_name}/{formatted_sim_name}"
36
+ default_experiment_name = (
37
+ f"/Users/{current_user.user_name}/{formatted_sim_name}"
38
+ )
37
39
 
38
40
  mlflow_experiment_name = (
39
41
  os.getenv("MLFLOW_EXPERIMENT_NAME") or default_experiment_name
@@ -103,8 +105,10 @@ def trace_risk_evaluation_fn(
103
105
  current_user = w.current_user.me()
104
106
 
105
107
  formatted_sim_name = simulation_name.lower().replace(" ", "_")
106
- default_experiment_name = f"/Users/{current_user.user_name}/{formatted_sim_name}"
107
-
108
+ default_experiment_name = (
109
+ f"/Users/{current_user.user_name}/{formatted_sim_name}"
110
+ )
111
+
108
112
  mlflow_experiment_name = (
109
113
  os.getenv("MLFLOW_EXPERIMENT_NAME") or default_experiment_name
110
114
  )
@@ -23,15 +23,18 @@ async def fetch_experiments(app_id: str = None) -> list[dict]:
23
23
  try:
24
24
  # get elapsed time for this request
25
25
  import time
26
+
26
27
  start_time = time.monotonic()
27
28
  experiments_response = await client.get(
28
- experiments_url,
29
- headers={"x-api-key": get_api_key_or_raise()},
30
- timeout=60.0, # Set timeout to 60 seconds
29
+ experiments_url,
30
+ headers={"x-api-key": get_api_key_or_raise()},
31
+ timeout=60.0, # Set timeout to 60 seconds
31
32
  )
32
33
  except httpx.ConnectTimeout:
33
34
  elapsed_time = time.monotonic() - start_time
34
- raise Exception(f"Warning: Connection timed out while fetching experiments. Elapsed time: {elapsed_time:.2f} seconds. Polling will continue. If this persists please contact Snowglobe support.")
35
+ raise Exception(
36
+ f"Warning: Connection timed out while fetching experiments. Elapsed time: {elapsed_time:.2f} seconds. Polling will continue. If this persists please contact Snowglobe support."
37
+ )
35
38
 
36
39
  if not experiments_response.status_code == 200:
37
40
  try:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: snowglobe
3
- Version: 0.4.11
3
+ Version: 0.4.13
4
4
  Summary: client server for usage with snowglobe experiments
5
5
  Author-email: Guardrails AI <contact@guardrailsai.com>
6
6
  License: MIT License
@@ -31,6 +31,7 @@ License-File: LICENSE
31
31
  Requires-Dist: typer>=0.9.0
32
32
  Requires-Dist: fastapi>=0.115.6
33
33
  Requires-Dist: uvicorn[standard]>=0.23.0
34
+ Requires-Dist: openai>=1.55.2
34
35
  Requires-Dist: requests>=2.31.0
35
36
  Requires-Dist: pydantic>=2.11.5
36
37
  Requires-Dist: APScheduler==4.0.0a6
@@ -87,7 +88,7 @@ import os
87
88
 
88
89
  client = OpenAI(api_key=os.getenv("SNOWGLOBE_API_KEY"))
89
90
 
90
- def process_scenario(request: CompletionRequest) -> CompletionFunctionOutputs:
91
+ def completion_fn(request: CompletionRequest) -> CompletionFunctionOutputs:
91
92
  """
92
93
  Process a scenario request from Snowglobe.
93
94
 
@@ -1,6 +1,7 @@
1
1
  typer>=0.9.0
2
2
  fastapi>=0.115.6
3
3
  uvicorn[standard]>=0.23.0
4
+ openai>=1.55.2
4
5
  requests>=2.31.0
5
6
  pydantic>=2.11.5
6
7
  APScheduler==4.0.0a6
@@ -141,11 +141,11 @@ class TestCli(unittest.TestCase):
141
141
  "snowglobe.client.src.cli.select_application_interactive"
142
142
  ) # Mock application selection in cli module
143
143
  @mock.patch(
144
- "snowglobe.client.src.cli.select_stateful_interactive"
145
- ) # Mock application selection in cli module
144
+ "snowglobe.client.src.cli.select_template_interactive"
145
+ ) # Mock template selection in cli module
146
146
  def test_init(
147
147
  self,
148
- mock_select_stateful,
148
+ mock_select_template,
149
149
  mock_select_app,
150
150
  mock_get_apps,
151
151
  mock_check_auth,
@@ -178,15 +178,15 @@ class TestCli(unittest.TestCase):
178
178
  ]
179
179
  mock_get_apps.return_value = (True, mock_applications, "Success")
180
180
 
181
- # Mock user selecting the first application
181
+ # Mock user selecting the first application and sync template
182
182
  mock_select_app.return_value = mock_applications[0]
183
- mock_select_stateful.return_value = False
183
+ mock_select_template.return_value = "sync"
184
184
  # Run the init command in a temporary directory
185
185
  with tempfile.TemporaryDirectory() as tmpdir:
186
186
  with mock.patch("os.getcwd", return_value=tmpdir):
187
187
  with mock.patch.dict(os.environ, {"SNOWGLOBE_API_KEY": "test_api_key"}):
188
188
  # Run the init command (pass default values for typer options)
189
- cli.init(file=None, stateful=False)
189
+ cli.init(file=None)
190
190
 
191
191
  # Verify authentication was checked
192
192
  mock_check_auth.assert_called_once()
@@ -240,9 +240,12 @@ class TestCli(unittest.TestCase):
240
240
  )
241
241
 
242
242
  # Check for required function definition
243
- self.assertIn(
244
- "def process_scenario(request: CompletionRequest) -> CompletionFunctionOutputs:",
245
- agent_content,
243
+ self.assertTrue(
244
+ "def completion(request: CompletionRequest) -> CompletionFunctionOutputs:"
245
+ in agent_content
246
+ or "def acompletion(request: CompletionRequest) -> CompletionFunctionOutputs:"
247
+ in agent_content,
248
+ "Agent file should contain either completion or acompletion function",
246
249
  )
247
250
 
248
251
  # Check that the file is valid Python (can be parsed)
@@ -264,10 +267,10 @@ class TestCli(unittest.TestCase):
264
267
  @mock.patch("snowglobe.client.src.cli.check_auth_status")
265
268
  @mock.patch("snowglobe.client.src.cli.get_remote_applications")
266
269
  @mock.patch("snowglobe.client.src.cli.select_application_interactive")
267
- @mock.patch("snowglobe.client.src.cli.select_stateful_interactive")
270
+ @mock.patch("snowglobe.client.src.cli.select_template_interactive")
268
271
  def test_init_with_file_option_creates_nested_path_and_mapping(
269
272
  self,
270
- mock_select_stateful,
273
+ mock_select_template,
271
274
  mock_select_app,
272
275
  mock_get_apps,
273
276
  mock_check_auth,
@@ -291,7 +294,7 @@ class TestCli(unittest.TestCase):
291
294
  ]
292
295
  mock_get_apps.return_value = (True, mock_applications, "Success")
293
296
  mock_select_app.return_value = mock_applications[0]
294
- mock_select_stateful.return_value = False # Default to stateless
297
+ mock_select_template.return_value = "sync" # Default to sync template
295
298
 
296
299
  runner = CliRunner()
297
300
 
@@ -460,7 +463,7 @@ class TestCli(unittest.TestCase):
460
463
  ) as mock_test_wrapper:
461
464
  mock_test_wrapper.return_value = (
462
465
  False,
463
- "Missing process_scenario function",
466
+ "Missing completion function",
464
467
  )
465
468
 
466
469
  with self.assertRaises(typer.Exit) as cm:
@@ -504,29 +507,29 @@ class TestCli(unittest.TestCase):
504
507
  self.assertFalse(result)
505
508
  # Should fail during module loading
506
509
 
507
- # Test Case 17: Missing process_scenario function
510
+ # Test Case 17: Missing completion function
508
511
  no_func_file = os.path.join(tmpdir, "no_func.py")
509
512
  with open(no_func_file, "w") as f:
510
- f.write("# Valid Python but no process_scenario function\npass\n")
513
+ f.write("# Valid Python but no completion function\npass\n")
511
514
 
512
515
  with mock.patch("os.getcwd", return_value=tmpdir):
513
516
  result, message = cli.test_agent_wrapper(
514
517
  "no_func.py", "test_123", "Test App"
515
518
  )
516
519
  self.assertFalse(result)
517
- self.assertIn("process_scenario function not found", message)
520
+ self.assertIn("completion or acompletion function not found", message)
518
521
 
519
- # Test Case 18: Non-callable process_scenario
522
+ # Test Case 18: Non-callable completion function
520
523
  non_callable_file = os.path.join(tmpdir, "non_callable.py")
521
524
  with open(non_callable_file, "w") as f:
522
- f.write("process_scenario = 'not a function'\n")
525
+ f.write("completion = 'not a function'\n")
523
526
 
524
527
  with mock.patch("os.getcwd", return_value=tmpdir):
525
528
  result, message = cli.test_agent_wrapper(
526
529
  "non_callable.py", "test_123", "Test App"
527
530
  )
528
531
  self.assertFalse(result)
529
- self.assertIn("process_scenario is not callable", message)
532
+ self.assertIn("completion function is not callable", message)
530
533
 
531
534
  # Test Case 19: Valid agent with default template response
532
535
  template_file = os.path.join(tmpdir, "template.py")
@@ -535,7 +538,7 @@ class TestCli(unittest.TestCase):
535
538
  """
536
539
  from snowglobe.client import CompletionRequest, CompletionFunctionOutputs
537
540
 
538
- def process_scenario(request):
541
+ def completion(request):
539
542
  return CompletionFunctionOutputs(response="Your response here")
540
543
  """
541
544
  )
File without changes
File without changes
File without changes