From f969098dc89c1ae454c91e129ca7d2e642766823 Mon Sep 17 00:00:00 2001
From: Joshua Smith <joshmsmith@gmail.com>
Date: Wed, 12 Mar 2025 09:55:25 -0400
Subject: [PATCH] finishing grok support

---
 activities/tool_activities.py | 44 +++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/activities/tool_activities.py b/activities/tool_activities.py
index 75e14e0..5d5ffc3 100644
--- a/activities/tool_activities.py
+++ b/activities/tool_activities.py
@@ -34,6 +34,7 @@ class ToolActivities:
 
         # Initialize client variables (all set to None initially)
         self.openai_client: Optional[OpenAI] = None
+        self.grok_client: Optional[OpenAI] = None
         self.anthropic_client: Optional[anthropic.Anthropic] = None
         self.genai_configured: bool = False
         self.deepseek_client: Optional[deepseek.DeepSeekAPI] = None
@@ -47,6 +48,13 @@ class ToolActivities:
                 print("Initialized OpenAI client")
             else:
                 print("Warning: OPENAI_API_KEY not set but LLM_PROVIDER is 'openai'")
+        
+        if self.llm_provider == "grok":
+            if os.environ.get("GROK_API_KEY"):
+                self.grok_client = OpenAI(api_key=os.environ.get("GROK_API_KEY"), base_url="https://api.x.ai/v1")
+                print("Initialized grok client")
+            else:
+                print("Warning: GROK_API_KEY not set but LLM_PROVIDER is 'grok'")
 
         elif self.llm_provider == "anthropic":
             if os.environ.get("ANTHROPIC_API_KEY"):
@@ -195,6 +203,8 @@ class ToolActivities:
             return self.prompt_llm_anthropic(input)
         elif self.llm_provider == "deepseek":
             return self.prompt_llm_deepseek(input)
+        elif self.llm_provider == "grok":
+            return self.prompt_llm_grok(input)
         else:
             return self.prompt_llm_openai(input)
 
@@ -244,6 +254,40 @@ class ToolActivities:
 
         return self.parse_json_response(response_content)
 
+    def prompt_llm_grok(self, input: ToolPromptInput) -> dict:
+        if not self.grok_client:
+            api_key = os.environ.get("GROK_API_KEY")
+            if not api_key:
+                raise ValueError(
+                    "GROK_API_KEY is not set in the environment variables but LLM_PROVIDER is 'grok'"
+                )
+            self.grok_client = OpenAI(api_key=api_key, base_url="https://api.x.ai/v1")
+            print("Initialized grok client on demand")
+
+        messages = [
+            {
+                "role": "system",
+                "content": input.context_instructions
+                + ". The current date is "
+                + datetime.now().strftime("%B %d, %Y"),
+            },
+            {
+                "role": "user",
+                "content": input.prompt,
+            },
+        ]
+
+        chat_completion = self.grok_client.chat.completions.create(
+            model="grok-2-1212", messages=messages  
+        )
+
+        response_content = chat_completion.choices[0].message.content
+        activity.logger.info(f"Grok response: {response_content}")
+
+        # Use the new sanitize function
+        response_content = self.sanitize_json_response(response_content)
+
+        return self.parse_json_response(response_content)
     def prompt_llm_ollama(self, input: ToolPromptInput) -> dict:
         # If not yet initialized, try to do so now (this is a backup if warm_up_ollama wasn't called or failed)
         if not self.ollama_initialized: