Spaces:
Running
Running
File size: 3,191 Bytes
6ff1f88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import httpx
import json
from typing import AsyncGenerator
async def run_model_stream(api_key: str, model: str, prompt: str):
"""
Run the SambaNova model with streaming response.
Args:
api_key: The API key to use for this request
model: The model name to use
prompt: The user's input prompt
Yields:
str: Chunks of the generated response
"""
try:
# Configure HTTP client with appropriate headers and SSL settings
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"Accept": "text/event-stream"
}
async with httpx.AsyncClient(
base_url="https://api.sambanova.ai/v1",
headers=headers,
verify=True,
timeout=httpx.Timeout(60.0, read=300.0)
) as client:
# Make streaming request
async with client.stream(
"POST",
"/chat/completions",
json={
"model": model, # Use model name directly from MODELS.csv
"messages": [
{
"role": "user",
"content": prompt
}
],
"stream": True,
"temperature": 0.7,
"max_tokens": 2048
}
) as response:
response.raise_for_status()
async for line in response.aiter_lines():
line = line.strip()
if not line:
continue
if line.startswith("data: "):
data = line[6:].strip()
if data == "[DONE]":
break
try:
chunk_data = json.loads(data)
if chunk_data.get("choices") and chunk_data["choices"][0].get("delta"):
content = chunk_data["choices"][0]["delta"].get("content")
if content:
yield content
except Exception as e:
print(f"Error parsing chunk: {e}")
continue
except httpx.HTTPError as e:
raise Exception(f"HTTP error with SambaNova API: {str(e)}")
except Exception as e:
raise Exception(f"Error with SambaNova API: {str(e)}")
async def run_model(api_key: str, model: str, prompt: str) -> str:
"""
Run the SambaNova model with the provided API key and prompt (non-streaming).
Args:
api_key: The API key to use for this request
model: The model name to use
prompt: The user's input prompt
Returns:
str: The generated response
"""
response_text = ""
async for chunk in run_model_stream(api_key, model, prompt):
response_text += chunk
return response_text |