fix: properly structure completions in Ollama provider

agentops/llms/providers/ollama.py

@@ -18,24 +18,30 @@
 
     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):
 
@@ -47,13 +53,16 @@
             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
 
@@ -96,21 +105,22 @@
             # 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