Skip to content

Commit

Permalink
fix: properly structure completions in Ollama provider (#589)
Browse files Browse the repository at this point in the history
* fix: properly structure completions in Ollama provider

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: apply ruff-format changes

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: add trailing commas and proper spacing

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: remove extra blank lines

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* docs: add Ollama example script

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: fix formatting in Ollama example script

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: remove extra blank lines from Ollama example script

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: remove trailing blank line from Ollama example script

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* style: apply ruff-format changes to ollama examples

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* fix: ensure proper event tracking in Ollama provider

Co-Authored-By: Alex Reibman <meta.alex.r@gmail.com>

* Delete examples/ollama_examples/ollama_examples.py

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Alex Reibman <meta.alex.r@gmail.com>
Co-authored-by: Alex Reibman <areibman@gmail.com>
  • Loading branch information
3 people authored Dec 17, 2024
1 parent 2a48bd6 commit c3fceb1
Showing 1 changed file with 19 additions and 9 deletions.
28 changes: 19 additions & 9 deletions agentops/llms/providers/ollama.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,24 +18,30 @@ class OllamaProvider(InstrumentedProvider):

def handle_response(self, response, kwargs, init_timestamp, session: Optional[Session] = None) -> dict:
llm_event = LLMEvent(init_timestamp=init_timestamp, params=kwargs)
if session is not None:
llm_event.session_id = session.session_id

def handle_stream_chunk(chunk: dict):
message = chunk.get("message", {"role": None, "content": ""})

if chunk.get("done"):
llm_event.completion["content"] += message.get("content")
llm_event.end_timestamp = get_ISO_time()
llm_event.model = f'ollama/{chunk.get("model")}'
llm_event.returns = chunk
llm_event.returns["message"] = llm_event.completion
llm_event.prompt = kwargs["messages"]
llm_event.agent_id = check_call_stack_for_agent_id()
self.client.record(llm_event)
self._safe_record(session, llm_event)

if llm_event.completion is None:
llm_event.completion = message
llm_event.completion = {
"role": message.get("role"),
"content": message.get("content", ""),
"tool_calls": None,
"function_call": None,
}
else:
llm_event.completion["content"] += message.get("content")
llm_event.completion["content"] += message.get("content", "")

if inspect.isgenerator(response):

Expand All @@ -47,13 +53,16 @@ def generator():
return generator()

llm_event.end_timestamp = get_ISO_time()

llm_event.model = f'ollama/{response["model"]}'
llm_event.returns = response
llm_event.agent_id = check_call_stack_for_agent_id()
llm_event.prompt = kwargs["messages"]
llm_event.completion = response["message"]

llm_event.completion = {
"role": response["message"].get("role"),
"content": response["message"].get("content", ""),
"tool_calls": None,
"function_call": None,
}
self._safe_record(session, llm_event)
return response

Expand Down Expand Up @@ -96,21 +105,22 @@ def patched_function(*args, **kwargs):
# Call the original function with its original arguments
init_timestamp = get_ISO_time()
result = original_func["ollama.Client.chat"](*args, **kwargs)
return self.handle_response(result, kwargs, init_timestamp)
return self.handle_response(result, kwargs, init_timestamp, session=kwargs.get("session", None))

# Override the original method with the patched one
Client.chat = patched_function

def _override_chat_async_client(self):
from ollama import AsyncClient

original_func = {}
original_func["ollama.AsyncClient.chat"] = AsyncClient.chat

async def patched_function(*args, **kwargs):
# Call the original function with its original arguments
init_timestamp = get_ISO_time()
result = await original_func["ollama.AsyncClient.chat"](*args, **kwargs)
return self.handle_response(result, kwargs, init_timestamp)
return self.handle_response(result, kwargs, init_timestamp, session=kwargs.get("session", None))

# Override the original method with the patched one
AsyncClient.chat = patched_function

0 comments on commit c3fceb1

Please sign in to comment.