From 00cf848a8fc212e0cdc0c2946bfd55d5f6d5ec9c Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Wed, 19 Mar 2025 15:04:22 -0700
Subject: [PATCH] fix tests

---
 crates/llm_gateway/src/stream_context.rs     |  9 ++++++++-
 crates/prompt_gateway/src/http_context.rs    |  2 +-
 demos/use_cases/orchestrating_agents/main.py | 12 +++++++++---
 tests/hurl/llm_gateway_simple.hurl           | 17 +++++++++++++++++
 4 files changed, 35 insertions(+), 5 deletions(-)
 create mode 100644 tests/hurl/llm_gateway_simple.hurl

diff --git a/crates/llm_gateway/src/stream_context.rs b/crates/llm_gateway/src/stream_context.rs
index ef00f8ec..b4a05575 100644
--- a/crates/llm_gateway/src/stream_context.rs
+++ b/crates/llm_gateway/src/stream_context.rs
@@ -504,6 +504,10 @@ impl HttpContext for StreamContext {
         };
 
         if self.streaming_response {
+            if body_utf8 == "data: [DONE]\n" {
+                return Action::Continue;
+            }
+
             let chat_completions_chunk_response_events =
                 match ChatCompletionStreamResponseServerEvents::try_from(body_utf8.as_str()) {
                     Ok(response) => response,
@@ -517,7 +521,10 @@ impl HttpContext for StreamContext {
                 };
 
             if chat_completions_chunk_response_events.events.is_empty() {
-                debug!("empty streaming response");
+                debug!(
+                    "cound't parse any streaming events: body str: {}",
+                    body_utf8
+                );
                 return Action::Continue;
             }
 
diff --git a/crates/prompt_gateway/src/http_context.rs b/crates/prompt_gateway/src/http_context.rs
index a26af5be..b8eb2797 100644
--- a/crates/prompt_gateway/src/http_context.rs
+++ b/crates/prompt_gateway/src/http_context.rs
@@ -40,7 +40,7 @@ impl HttpContext for StreamContext {
                     if endpoints.len() == 1 {
                         let (name, _) = endpoints.iter().next().unwrap();
                         debug!("Setting ARCH_PROVIDER_HINT_HEADER to {}", name);
-                        self.set_http_request_header(ARCH_ROUTING_HEADER, Some(&name));
+                        self.set_http_request_header(ARCH_ROUTING_HEADER, Some(name));
                     } else {
                         warn!("Need single endpoint when use_agent_orchestrator is set");
                         self.send_server_error(
diff --git a/demos/use_cases/orchestrating_agents/main.py b/demos/use_cases/orchestrating_agents/main.py
index 4c25314a..27a9598a 100644
--- a/demos/use_cases/orchestrating_agents/main.py
+++ b/demos/use_cases/orchestrating_agents/main.py
@@ -26,13 +26,14 @@ class ChatCompletionsRequest(BaseModel):
 
 
 openai_client = openai.OpenAI(
-    api_key=" -- ",  # Replace with your OpenAI API key
+    api_key="None",  # archgw picks the API key from the config file
+    base_url="http://host.docker.internal:12000/v1",
 )
 
 
 def call_openai(messages: List[Dict[str, str]], stream: bool):
     completion = openai_client.chat.completions.create(
-        model="gpt-4o",  # archgw picks the default LLM configured in the config file
+        model="None",  # archgw picks the default LLM configured in the config file
         messages=messages,
         stream=stream,
     )
@@ -93,7 +94,12 @@ AGENTS = {
 @app.post("/v1/chat/completions")
 def completion_api(req: ChatCompletionsRequest, request: Request):
     agent_name = req.metadata.get("agent-name", "unknown_agent")
-    agent = AGENTS.get(agent_name, AGENTS["unknown_agent"])
+    agent = AGENTS.get(agent_name)
     logger.info(f"Routing to agent: {agent_name}")
 
     return agent.handle(req)
+
+
+@app.get("/healthz")
+async def healthz():
+    return {"status": "ok"}
diff --git a/tests/hurl/llm_gateway_simple.hurl b/tests/hurl/llm_gateway_simple.hurl
new file mode 100644
index 00000000..f16e0074
--- /dev/null
+++ b/tests/hurl/llm_gateway_simple.hurl
@@ -0,0 +1,17 @@
+POST http://localhost:12000/v1/chat/completions
+Content-Type: application/json
+
+{
+  "model": "--",
+  "messages": [
+    {
+      "role": "system",
+      "content": "You are a helpful assistant.\n"
+    },
+    {
+      "role": "user",
+      "content": "I want to sell red shoes"
+    }
+  ],
+  "stream": true
+}