add handle streming response for anthropic

This commit is contained in:
Zishan committed 2025-02-11 15:39:22 +01:00
1 parent 87fc58e4de
commit b832ca2cf1
2 files changed
+152 -5

No files matched your search

@@ -0,0 +1,51 @@
import anthropic
from typing import Dict
import os
from tavily import TavilyClient
import anthropic
from httpx import Client
import os
# from invariant import testing
tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
import datetime
def test_streaming_response_without_toolcall():
# Example queries
dataset_name = "claude_streaming_agent_test" + str(datetime.datetime.now().strftime("%Y%m%d%H%M%S"))
invariant_api_key = os.environ.get("INVARIANT_API_KEY")
client = anthropic.Anthropic(
http_client=Client(
headers={
"Invariant-Authorization": f"Bearer {invariant_api_key}"
},
),
base_url=f"http://localhost/api/v1/proxy/{dataset_name}/anthropic",
)
cities = ["Zurich", "New York", "London"]
queries = [
"Can you introduce Zurich city within 200 words?",
"Tell me the history of New York",
"How's the weather in London next week?"
]
# Process each query
for index,query in enumerate(queries):
messages = [
{
"role": "user",
"content": query
}
]
response_text = ""
with client.messages.stream(
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=messages,
# stream = True
) as response:
for reply in response.text_stream:
response_text += reply
print(reply, end="", flush=True)
assert reply != ""
assert cities[index] in response_text