Merge branch 'main' into kt/add-af-cli

log10/_httpx_utils.py

@@ -9,7 +9,7 @@
 from httpx import Request, Response
 
 from log10.llm import Log10Config
-from log10.load import get_log10_session_tags, session_id_var
+from log10.load import get_log10_session_tags, last_completion_response_var, session_id_var
 
 
 logger: logging.Logger = logging.getLogger("LOG10")
@@ -162,6 +162,7 @@ async def log_request(request: Request):
  if not completion_id:
  return
 
+ last_completion_response_var.set({"completionID": completion_id})
  orig_module = ""
  orig_qualname = ""
  request_content_decode = request.content.decode("utf-8")

log10/litellm.py

@@ -2,6 +2,7 @@
 from typing import List, Optional
 
 from log10.llm import LLM, Kind, Log10Config
+from log10.load import last_completion_response_var
 
 
 try:
@@ -56,6 +57,7 @@ def log_pre_api_call(self, model, messages, kwargs):
  request = kwargs.get("additional_args").get("complete_input_dict").copy()
  request["messages"] = messages.copy()
  completion_id = self.log_start(request, Kind.chat, self.tags)
+ last_completion_response_var.set({"completionID": completion_id})
  litellm_call_id = kwargs.get("litellm_call_id")
  self.runs[litellm_call_id] = {
  "kind": Kind.chat,
@@ -76,7 +78,10 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time):
  litellm_call_id = kwargs.get("litellm_call_id")
  run = self.runs.get(litellm_call_id, None)
  duration = (end_time - start_time).total_seconds()
- self.log_end(run["completion_id"], response_obj.dict(), duration)
+
+ completion_id = run["completion_id"]
+ last_completion_response_var.set({"completionID": completion_id})
+ self.log_end(completion_id, response_obj.dict(), duration)
 
  def log_failure_event(self, kwargs, response_obj, start_time, end_time):
  update_log_row = {

log10/load.py

@@ -153,6 +153,10 @@ def last_completion_url(self):
  if last_completion_response_var.get() is None:
  return None
  response = last_completion_response_var.get()
+
+ # organizationSlug will not be returned from httpx hook
+ if not response.get("organizationSlug", ""):
+ return None
  return f'{url}/app/{response["organizationSlug"]}/completions/{response["completionID"]}'
 
  def last_completion_id(self):
@@ -186,8 +190,8 @@ async def log_async(completion_url, log_row):
  res = None
  try:
  res = post_request(completion_url)
- last_completion_response_var.set(res.json())
  completionID = res.json().get("completionID", None)
+ organizationSlug = res.json().get("organizationSlug", None)
 
  if completionID is None:
  logging.warn("LOG10: failed to get completionID from log10. Skipping log.")
@@ -212,7 +216,7 @@ async def log_async(completion_url, log_row):
  logging.warn(f"LOG10: failed to log: {e}. Skipping")
  return None
 
- return completionID
+ return {"completionID": completionID, "organizationSlug": organizationSlug}
 
 
 def run_async_in_thread(completion_url, log_row, result_queue):
@@ -671,7 +675,9 @@ def wrapper(*args, **kwargs):
  with timed_block("extra time spent waiting for log10 call"):
  while result_queue.empty():
  pass
- completionID = result_queue.get()
+ result = result_queue.get()
+ completionID = result["completionID"]
+ last_completion_response_var.set(result)
 
  if completionID is None:
  logger.warning(f"LOG10: failed to get completionID from log10: {e}. Skipping log.")
@@ -698,7 +704,9 @@ def wrapper(*args, **kwargs):
  with timed_block("extra time spent waiting for log10 call"):
  while result_queue.empty():
  pass
- completionID = result_queue.get()
+ result = result_queue.get()
+ completionID = result["completionID"]
+ last_completion_response_var.set(result)
 
  with timed_block("result call duration (sync)"):
  response = output

tests/conftest.py

@@ -1,5 +1,7 @@
 import pytest
 
+from log10.load import log10_session
+
 
 def pytest_addoption(parser):
  parser.addoption("--openai_model", action="store", help="Model name for OpenAI tests")
@@ -59,3 +61,10 @@ def google_model(request):
 @pytest.fixture
 def magentic_model(request):
  return request.config.getoption("--magentic_model")
+
+
+@pytest.fixture
+def session():
+ with log10_session() as session:
+ assert session.last_completion_id() is None, "No completion ID should be found."
+ yield session

tests/test_anthropic.py

@@ -1,4 +1,5 @@
 import base64
+import json
 
 import anthropic
 import httpx
@@ -8,13 +9,14 @@
 from typing_extensions import override
 
 from log10.load import log10
+from tests.utils import _LogAssertion
 
 
 log10(anthropic)
 
 
 @pytest.mark.chat
-def test_messages_create(anthropic_model):
+def test_messages_create(session, anthropic_model):
  client = anthropic.Anthropic()
 
  message = client.messages.create(
@@ -27,13 +29,13 @@ def test_messages_create(anthropic_model):
 
  text = message.content[0].text
  assert isinstance(text, str)
- assert text, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()
 
 
 @pytest.mark.chat
 @pytest.mark.async_client
 @pytest.mark.asyncio
-async def test_messages_create_async(anthropic_model):
+async def test_messages_create_async(session, anthropic_model):
  client = anthropic.AsyncAnthropic()
 
  message = await client.messages.create(
@@ -46,12 +48,12 @@ async def test_messages_create_async(anthropic_model):
 
  text = message.content[0].text
  assert isinstance(text, str)
- assert text, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()
 
 
 @pytest.mark.chat
 @pytest.mark.stream
-def test_messages_create_stream(anthropic_model):
+def test_messages_create_stream(session, anthropic_model):
  client = anthropic.Anthropic()
 
  stream = client.messages.create(
@@ -75,11 +77,11 @@ def test_messages_create_stream(anthropic_model):
  if text.isdigit():
  assert int(text) <= 10
 
- assert output, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=output).assert_chat_response()
 
 
 @pytest.mark.vision
-def test_messages_image(anthropic_model):
+def test_messages_image(session, anthropic_model):
  client = anthropic.Anthropic()
 
  image1_url = "https://upload.wikimedia.org/wikipedia/commons/a/a7/Camponotus_flavomarginatus_ant.jpg"
@@ -108,12 +110,11 @@ def test_messages_image(anthropic_model):
  )
 
  text = message.content[0].text
- assert text, "No output from the model."
- assert "ant" in text
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()
 
 
 @pytest.mark.chat
-def test_chat_not_given(anthropic_model):
+def test_chat_not_given(session, anthropic_model):
  client = anthropic.Anthropic()
 
  message = client.beta.tools.messages.create(
@@ -132,10 +133,11 @@ def test_chat_not_given(anthropic_model):
  content = message.content[0].text
  assert isinstance(content, str)
  assert content, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=content).assert_chat_response()
 
 
 @pytest.mark.chat
-def test_beta_tools_messages_create(anthropic_model):
+def test_beta_tools_messages_create(session, anthropic_model):
  client = anthropic.Anthropic()
 
  message = client.beta.tools.messages.create(
@@ -145,13 +147,13 @@ def test_beta_tools_messages_create(anthropic_model):
  )
 
  text = message.content[0].text
- assert text, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()
 
 
 @pytest.mark.chat
 @pytest.mark.async_client
 @pytest.mark.asyncio
-async def test_beta_tools_messages_create_async(anthropic_model):
+async def test_beta_tools_messages_create_async(session, anthropic_model):
  client = anthropic.AsyncAnthropic()
 
  message = await client.beta.tools.messages.create(
@@ -161,13 +163,13 @@ async def test_beta_tools_messages_create_async(anthropic_model):
  )
 
  text = message.content[0].text
- assert text, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()
 
 
 @pytest.mark.chat
 @pytest.mark.stream
 @pytest.mark.context_manager
-def test_messages_stream_context_manager(anthropic_model):
+def test_messages_stream_context_manager(session, anthropic_model):
  client = anthropic.Anthropic()
 
  output = ""
@@ -187,15 +189,15 @@ def test_messages_stream_context_manager(anthropic_model):
  if hasattr(message.delta, "text"):
  output += message.delta.text
 
- assert output, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=output).assert_chat_response()
 
 
 @pytest.mark.chat
 @pytest.mark.stream
 @pytest.mark.context_manager
 @pytest.mark.async_client
 @pytest.mark.asyncio
-async def test_messages_stream_context_manager_async(anthropic_model):
+async def test_messages_stream_context_manager_async(session, anthropic_model):
  client = anthropic.AsyncAnthropic()
 
  output = ""
@@ -212,13 +214,13 @@ async def test_messages_stream_context_manager_async(anthropic_model):
  async for text in stream.text_stream:
  output += text
 
- assert output, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=output).assert_chat_response()
 
 
 @pytest.mark.tools
 @pytest.mark.stream
 @pytest.mark.context_manager
-def test_tools_messages_stream_context_manager(anthropic_model):
+def test_tools_messages_stream_context_manager(session, anthropic_model):
  client = anthropic.Anthropic()
  output = ""
  with client.beta.tools.messages.stream(
@@ -252,23 +254,25 @@ def test_tools_messages_stream_context_manager(anthropic_model):
  if hasattr(message.delta, "partial_json"):
  output += message.delta.partial_json
 
- assert output, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=output).assert_chat_response()
 
 
 @pytest.mark.tools
 @pytest.mark.stream
 @pytest.mark.context_manager
 @pytest.mark.async_client
 @pytest.mark.asyncio
-async def test_tools_messages_stream_context_manager_async(anthropic_model):
+async def test_tools_messages_stream_context_manager_async(session, anthropic_model):
  client = anthropic.AsyncAnthropic()
- output = None
+ json_snapshot = None
+ final_message = None
+ output = ""
 
  class MyHandler(AsyncToolsBetaMessageStream):
  @override
  async def on_input_json(self, delta: str, snapshot: object) -> None:
- nonlocal output
- output = snapshot
+ nonlocal json_snapshot
+ json_snapshot = snapshot
 
  async with client.beta.tools.messages.stream(
  model=anthropic_model,
@@ -294,6 +298,16 @@ async def on_input_json(self, delta: str, snapshot: object) -> None:
  max_tokens=1024,
  event_handler=MyHandler,
  ) as stream:
- await stream.until_done()
+ final_message = await stream.get_final_message()
+
+ content = final_message.content[0]
+ if hasattr(content, "text"):
+ output = content.text
+
+ if json_snapshot:
+ output += json.dumps(json_snapshot)
 
  assert output, "No output from the model."
+ assert session.last_completion_id(), "No completion ID found."
+ ## TODO fix this test after the anthropic fixes for the tool_calls
+ # _LogAssertion(completion_id=session.last_completion_id(), message_content=output).assert_chat_response()

tests/test_google.py

@@ -2,13 +2,14 @@
 import pytest
 
 from log10.load import log10
+from tests.utils import _LogAssertion
 
 
 log10(genai)
 
 
 @pytest.mark.chat
-def test_genai_chat(google_model):
+def test_genai_chat(session, google_model):
  model = genai.GenerativeModel(google_model)
  chat = model.start_chat()
 
@@ -21,11 +22,11 @@ def test_genai_chat(google_model):
 
  text = response.text
  assert isinstance(text, str)
- assert text, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()
 
 
 @pytest.mark.chat
-def test_genai_chat_w_history(google_model):
+def test_genai_chat_w_history(session, google_model):
  model = genai.GenerativeModel(google_model, system_instruction="You are a cat. Your name is Neko.")
  chat = model.start_chat(
  history=[
@@ -39,4 +40,4 @@ def test_genai_chat_w_history(google_model):
 
  text = response.text
  assert isinstance(text, str)
- assert text, "No output from the model."
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=text).assert_chat_response()

tests/test_lamini.py

@@ -2,15 +2,17 @@
 import pytest
 
 from log10.load import log10
+from tests.utils import _LogAssertion
 
 
 log10(lamini)
 
 
 @pytest.mark.chat
-def test_generate(lamini_model):
+def test_generate(session, lamini_model):
  llm = lamini.Lamini(lamini_model)
  response = llm.generate("What's 2 + 9 * 3?")
 
  assert isinstance(response, str)
  assert "29" in response
+ _LogAssertion(completion_id=session.last_completion_id(), message_content=response).assert_chat_response()