From 28407f7a56597cd289be03b2c965284eb3e5130e Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.org>
Date: Mon, 28 Oct 2024 19:18:05 -0400
Subject: [PATCH] whoot

---
 README.md                        |  5 +-
 simplemind/providers/__init__.py |  3 +-
 simplemind/providers/xai.py      | 80 ++++++++++++++++++++++++++++++++
 simplemind/settings.py           |  1 +
 t.py                             | 16 +++----
 5 files changed, 92 insertions(+), 13 deletions(-)
 create mode 100644 simplemind/providers/xai.py

diff --git a/README.md b/README.md
index bf37745..a62a18f 100644
--- a/README.md
+++ b/README.md
@@ -72,8 +72,9 @@ SimpleMind also allows for easy conversational flows:
 ```
 
 ## Supported APIs
-- **OpenAI GPT**
-= **Anthropic**
+- **OpenAI's GPT**
+- **Anthropic's Claude**
+- **X.AI's Grok**
 
 ## Why SimpleMind?
 - **Intuitive**: Built with Pythonic simplicity and readability in mind.
diff --git a/simplemind/providers/__init__.py b/simplemind/providers/__init__.py
index 3f1ed4a..87787ea 100644
--- a/simplemind/providers/__init__.py
+++ b/simplemind/providers/__init__.py
@@ -1,4 +1,5 @@
 from .openai import OpenAI
 from .anthropic import Anthropic
+from .xai import XAI
 
-providers = [OpenAI, Anthropic]
+providers = [OpenAI, Anthropic, XAI]
diff --git a/simplemind/providers/xai.py b/simplemind/providers/xai.py
new file mode 100644
index 0000000..f966105
--- /dev/null
+++ b/simplemind/providers/xai.py
@@ -0,0 +1,80 @@
+import openai as oa
+import instructor
+
+# from ..models import Conversation, Message
+from ..settings import settings
+
+PROVIDER_NAME = "xai"
+DEFAULT_MODEL = "grok-beta"
+BASE_URL = "https://api.x.ai/v1"
+DEFAULT_MAX_TOKENS = 1000
+
+
+class XAI:
+    __name__ = PROVIDER_NAME
+    DEFAULT_MODEL = DEFAULT_MODEL
+
+    def __init__(self, api_key: str = None):
+        self.api_key = api_key or settings.XAI_API_KEY
+
+    @property
+    def client(self):
+        """The raw OpenAI client."""
+
+        return oa.OpenAI(
+            api_key=settings.XAI_API_KEY,
+            base_url="https://api.x.ai/v1",
+        )
+
+    @property
+    def structured_client(self):
+        """A client patched with Instructor."""
+        return instructor.patch(
+            oa.OpenAI(api_key=self.api_key, base_url="https://api.x.ai/v1")
+        )
+
+    def send_conversation(self, conversation: "Conversation"):
+        """Send a conversation to the OpenAI API."""
+        from ..models import Message
+
+        messages = [
+            {"role": msg.role, "content": msg.text} for msg in conversation.messages
+        ]
+
+        response = self.client.chat.completions.create(
+            model=conversation.llm_model or DEFAULT_MODEL, messages=messages
+        )
+
+        # Get the response content from the OpenAI response
+        assistant_message = response.choices[0].message
+
+        # Create and return a properly formatted Message instance
+        return Message(
+            role="assistant",
+            text=assistant_message.content,
+            raw=response,
+            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_provider=PROVIDER_NAME,
+        )
+
+    def structured_response(self, prompt, response_model, *, llm_model):
+        # Ensure messages are provided in kwargs
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.structured_client.chat.completions.create(
+            messages=messages, model=llm_model, response_model=response_model
+        )
+        return response
+
+    def generate_text(self, prompt, *, llm_model):
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.chat.completions.create(
+            messages=messages, model=llm_model
+        )
+
+        return response.choices[0].message.content
diff --git a/simplemind/settings.py b/simplemind/settings.py
index 0189beb..7828508 100644
--- a/simplemind/settings.py
+++ b/simplemind/settings.py
@@ -5,6 +5,7 @@ from pydantic_settings import BaseSettings
 class Settings(BaseSettings):
     OPENAI_API_KEY: str = Field(..., env="OPENAI_API_KEY")
     ANTHROPIC_API_KEY: str = Field(..., env="ANTHROPIC_API_KEY")
+    XAI_API_KEY: str = Field(..., env="XAI_API_KEY")
 
 
 settings = Settings()
diff --git a/t.py b/t.py
index 4f431d5..4bb98ab 100644
--- a/t.py
+++ b/t.py
@@ -2,16 +2,12 @@ import simplemind as sm
 from pydantic import BaseModel
 
 
-class Poem(BaseModel):
-    title: str
-    content: str
+conversation = sm.create_conversation(llm_model="grok-beta", llm_provider="xai")
 
-
-output = sm.generate_data(
-    "Write a poem about love",
-    llm_model="gpt-4o-mini",
-    llm_provider="openai",
-    response_model=Poem,
+conversation.add_message(
+    role="user",
+    text="Write a poem about love",
 )
+r = conversation.send()
 
-print(output)
+print(r)