58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
from typing import List, Union, Generator, Iterator
|
|
from pydantic import BaseModel
|
|
import logging
|
|
import os
|
|
import requests
|
|
import json
|
|
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
class Pipeline:
|
|
class Valves(BaseModel):
|
|
OR_MODEL: str
|
|
OR_URL: str
|
|
OR_KEY: str
|
|
|
|
def __init__(self):
|
|
self.name = "Basic Openrouter Chat"
|
|
self.valves = self.Valves(
|
|
**{
|
|
"pipelines": ["*"],
|
|
"OR_MODEL": os.getenv("OR_MODEL", "anthropic/claude-3.5-haiku:beta"),
|
|
"OR_URL": os.getenv("OR_URL", "https://openrouter.ai/api/v1/chat/completions"),
|
|
"OR_KEY": os.getenv("OR_KEY", "OPENROUTER_API_KEY")
|
|
}
|
|
)
|
|
|
|
async def on_startup(self):
|
|
pass
|
|
|
|
async def on_shutdown(self):
|
|
pass
|
|
|
|
def pipe(self, user_message: str, model_id: str, messages: List[dict], body: dict) -> Union[str, Generator, Iterator]:
|
|
try:
|
|
response = requests.post(
|
|
url = self.valves.OR_URL,
|
|
headers = {
|
|
"Authorization": f"Bearer {self.valves.OR_KEY}"
|
|
},
|
|
|
|
data = json.dumps({
|
|
"model": self.valves.OR_MODEL,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": user_message
|
|
}
|
|
]
|
|
})
|
|
)
|
|
|
|
return response
|
|
except requests.HTTPError as e:
|
|
logging.error(f"Clientresponse error: {e}")
|
|
return "HTTP backend error"
|
|
except Exception as e:
|
|
logging.error(f"Unexpected error: {e}")
|
|
return f"Unexpected error: {e}" |