Skip to content

feat: add xai support #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
PavelInjective wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

feat: add xai support #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions agent_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,23 @@ def __init__(self):
raise ValueError(
"No OpenAI API key found. Please set the OPENAI_API_KEY environment variable."
)


self.provider = os.getenv("API_PROVIDER", "openai")
match self.provider:
case "openai":
self.initial_model = "gpt-4o"
self.conversational_model = "gpt-4-turbo-preview"
self.base_url = "https://api.openai.com/v1"
case "xai":
self.initial_model = "grok-2-latest"
self.conversational_model = "grok-2-latest"
self.base_url = "https://api.x.ai/v1"

# Initialize OpenAI client
self.client = OpenAI(api_key=self.api_key)
self.client = OpenAI(
api_key=self.api_key,
base_url=self.base_url,
)

# Initialize conversation histories
self.conversations = {}
Expand Down Expand Up @@ -107,7 +121,7 @@ async def get_response(
# Get response from OpenAI
response = await asyncio.to_thread(
self.client.chat.completions.create,
model="gpt-4o",
model=self.initial_model,
messages=[
{
"role": "system",
Expand Down Expand Up @@ -183,7 +197,7 @@ async def get_response(
# Get final response
second_response = await asyncio.to_thread(
self.client.chat.completions.create,
model="gpt-4-turbo-preview",
model=self.conversational_model,
messages=self.conversations[session_id],
max_tokens=2000,
temperature=0.7,
Expand Down
3 changes: 2 additions & 1 deletion readme.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,10 @@ All required packages are listed in requirements.txt. Ensure you have Python 3.1
```bash
pip install -r requirements.txt
```
3. Setup OpenAI api Key
3. Setup OpenAI api Key and API provider (`"openai"` is set by default; `"xai"` is also supported).
```bash
export OPENAI_API_KEY="your_openai_api_key_here"
export API_PROVIDER="openai|xai"
```
4. Running the agent
To start the backend on a specified port (default is 5000), run:
Expand Down