From 98bed2a24837756100c4518220f912d34680d3c1 Mon Sep 17 00:00:00 2001
From: maamalama <bokarevstefan475@gmail.com>
Date: Wed, 24 Jul 2024 17:40:20 -0700
Subject: [PATCH] tools_call to Helicone

---
 litellm/integrations/helicone.py | 59 ++++++++++++++++++++++----------
 1 file changed, 41 insertions(+), 18 deletions(-)

diff --git a/litellm/integrations/helicone.py b/litellm/integrations/helicone.py
index f0666431a69..edab68c1c68 100644
--- a/litellm/integrations/helicone.py
+++ b/litellm/integrations/helicone.py
@@ -17,27 +17,50 @@ class HeliconeLogger:
         self.key = os.getenv("HELICONE_API_KEY")
 
     def claude_mapping(self, model, messages, response_obj):
-        from anthropic import HUMAN_PROMPT, AI_PROMPT
+      from anthropic import HUMAN_PROMPT, AI_PROMPT
 
-        prompt = f"{HUMAN_PROMPT}"
-        for message in messages:
-            if "role" in message:
-                if message["role"] == "user":
-                    prompt += f"{HUMAN_PROMPT}{message['content']}"
-                else:
-                    prompt += f"{AI_PROMPT}{message['content']}"
-            else:
-                prompt += f"{HUMAN_PROMPT}{message['content']}"
-        prompt += f"{AI_PROMPT}"
-        claude_provider_request = {"model": model, "prompt": prompt}
+      prompt = f"{HUMAN_PROMPT}"
+      for message in messages:
+          if "role" in message:
+              if message["role"] == "user":
+                  prompt += f"{HUMAN_PROMPT}{message['content']}"
+              else:
+                  prompt += f"{AI_PROMPT}{message['content']}"
+          else:
+              prompt += f"{HUMAN_PROMPT}{message['content']}"
+      prompt += f"{AI_PROMPT}"
+      claude_provider_request = {"model": model, "prompt": prompt}
 
-        claude_response_obj = {
-            "completion": response_obj["choices"][0]["message"]["content"],
-            "model": model,
-            "stop_reason": "stop_sequence",
-        }
+      choice = response_obj["choices"][0]
+      message = choice["message"]
 
-        return claude_provider_request, claude_response_obj
+      content = []
+      if "tool_calls" in message and message["tool_calls"]:
+          for tool_call in message["tool_calls"]:
+              content.append({
+                  "type": "tool_use",
+                  "id": tool_call["id"],
+                  "name": tool_call["function"]["name"],
+                  "input": tool_call["function"]["arguments"]
+              })
+      elif "content" in message and message["content"]:
+          content = [{"type": "text", "text": message["content"]}]
+
+      claude_response_obj = {
+          "id": response_obj["id"],
+          "type": "message",
+          "role": "assistant",
+          "model": model,
+          "content": content,
+          "stop_reason": choice["finish_reason"],
+          "stop_sequence": None,
+          "usage": {
+              "input_tokens": response_obj["usage"]["prompt_tokens"],
+              "output_tokens": response_obj["usage"]["completion_tokens"]
+          }
+      }
+
+      return claude_provider_request, claude_response_obj
     
     @staticmethod
     def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict: