feat: OpenRouter engine, with OPENROUTER_API_KEY env var
This commit is contained in:
1 parent
4e15b44d39
commit
50cd677515
7 files changed
+58
-119
No files matched your search
@@ -7,4 +7,5 @@ replicate
|
||||
qdrant-client[fastembed]
|
||||
ollama
|
||||
tortoise-orm
|
||||
llama-cpp-python
|
||||
llama-cpp-python
|
||||
openrouter
|
||||
@@ -3,8 +3,6 @@ import json
|
||||
import inspect
|
||||
import time
|
||||
|
||||
from inference import LlamaMessage, replicate_llama_completion
|
||||
import gral
|
||||
from temporary import generate_tmp_id
|
||||
|
||||
# Tool related classes must be moved to a different module
|
||||
|
||||
@@ -3,7 +3,7 @@ import json
|
||||
import inspect
|
||||
from string import Template
|
||||
|
||||
from inference import LlamaMessage, replicate_llama_completion
|
||||
from inference import LlamaMessage, replicate_completion
|
||||
from agent import Agent, AgentField
|
||||
|
||||
# Info and defaults
|
||||
@@ -137,7 +137,7 @@ class CharacterAgent(Agent):
|
||||
LlamaMessage("system", self.system_prompt),
|
||||
LlamaMessage("user", content)
|
||||
]
|
||||
response = await replicate_llama_completion(self.app, messages)
|
||||
response = await replicate_completion(self.app, messages)
|
||||
return response
|
||||
|
||||
async def received_message(self, message):
|
||||
@@ -201,7 +201,7 @@ class CharacterAgent(Agent):
|
||||
# TODO: Inject system prompt before the last user message
|
||||
# If not any user message, combine with first system prompt
|
||||
|
||||
controller_response = await replicate_llama_completion(self.app, llama_messages)
|
||||
controller_response = await replicate_completion(self.app, llama_messages)
|
||||
|
||||
if self.debug:
|
||||
await self.app.send_chat_message(
|
||||
@@ -282,7 +282,7 @@ class CharacterAgent(Agent):
|
||||
memories = [it.document for it in query_result]
|
||||
|
||||
print("rerank memories")
|
||||
rerank_response = await replicate_llama_completion(self.app, [
|
||||
rerank_response = await replicate_completion(self.app, [
|
||||
LlamaMessage("system", rerank_system_prompt),
|
||||
LlamaMessage("user", rerank_prompt_template.substitute(
|
||||
conversation=chat_history,
|
||||
|
||||
@@ -3,7 +3,6 @@ import json
|
||||
import inspect
|
||||
from string import Template
|
||||
|
||||
from inference import LlamaMessage, replicate_llama_completion, Models
|
||||
from agent import Agent, AgentField
|
||||
|
||||
CLASS_DESCRIPTION = """
|
||||
|
||||
+8
-1
@@ -10,6 +10,7 @@ from tortoise.expressions import Q
|
||||
from quart import Quart, Response, request, send_from_directory, jsonify, websocket
|
||||
from qdrant_client import QdrantClient
|
||||
import replicate
|
||||
from openrouter import OpenRouter
|
||||
|
||||
import database as db
|
||||
import agent as agent_mod
|
||||
@@ -36,7 +37,13 @@ class App(ChatMixin):
|
||||
if replicate_api_token:
|
||||
self.replicate_client = replicate.Client(api_token=replicate_api_token)
|
||||
else:
|
||||
raise Exception("Horde only works with a replicate connection. Please provide an API Token.")
|
||||
self.replicate_client = None
|
||||
|
||||
openrouter_api_key = os.getenv("OPENROUTER_API_KEY")
|
||||
if openrouter_api_key:
|
||||
self.openrouter_client = OpenRouter(api_key=openrouter_api_key)
|
||||
else:
|
||||
self.openrouter_client = None
|
||||
|
||||
self.agent_classes = [
|
||||
agents_mod.ExpertAgent,
|
||||
|
||||
@@ -1,76 +0,0 @@
|
||||
from inference import LlamaMessage, replicate_llama_completion, Models
|
||||
import re
|
||||
|
||||
message_format_section = """
|
||||
When suggesting code to the user, keep the code as small snippets.
|
||||
You have the capability to update files in your current environment with file tags before your code suggestions.
|
||||
If the update is meant to be applied on a file, include a file tag line immediately before the code block like this: "[filename.txt]
|
||||
```
|
||||
updated segment
|
||||
```"
|
||||
Omit the file link when the block is not related to a known file,
|
||||
Code blocks must be under 30 lines. Your responses MUST NOT contain horizontal rules.
|
||||
When suggesting bash code meant to be directly executed, use the same syntax as above but marking the code block with "[eval]" instead.
|
||||
"""
|
||||
|
||||
|
||||
async def gral_message(app, system_prompt, message_history: list, file_map: dict):
|
||||
messages = [
|
||||
LlamaMessage("system", "\n".join([
|
||||
system_prompt,
|
||||
message_format_section,
|
||||
])),
|
||||
]
|
||||
|
||||
file_parts = []
|
||||
for filename, content in file_map.items():
|
||||
file_parts.append(f"\n---\n[{filename}]\n\n{content}\n")
|
||||
|
||||
if file_parts:
|
||||
file_lines = [
|
||||
"This message contains all files relevant to the prompt: each file in this message is separated with horizontal rules \"---\", followed by the file name in square brackets, followed by the file content.",
|
||||
"Only this message contains horizontal rules, your responses never has them.",
|
||||
"Conversation start in the message following this one."
|
||||
]
|
||||
file_lines.extend(file_parts)
|
||||
messages.append(
|
||||
LlamaMessage("user", "\n".join(file_lines))
|
||||
)
|
||||
|
||||
messages.extend([LlamaMessage(msg['role'], msg['content']) for msg in message_history])
|
||||
|
||||
response = await replicate_llama_completion(app, messages, Models.llama_70b)
|
||||
return response
|
||||
|
||||
async def gral_update_file(app, system_prompt, message_history: list, filepath: str, file_content: str):
|
||||
messages = [
|
||||
LlamaMessage("system", "\n".join([
|
||||
system_prompt,
|
||||
"Respond with the raw contents of the file updated as concluded in the conversation.",
|
||||
])),
|
||||
]
|
||||
|
||||
# History Message
|
||||
message_lines = [f"This is the history of the conversation that led to the update of the file {filepath}."]
|
||||
message_lines.extend([f"{entry['role']}: {entry['content']}" for entry in message_history if entry.get('content')])
|
||||
messages.append(LlamaMessage("user", "\n".join(message_lines)))
|
||||
|
||||
# Main Instruction Message
|
||||
message_lines = [
|
||||
f"Write the updated contents of the file: {filepath}, as concluded in the conversation above.",
|
||||
f"The current content of the file {filepath} is pasted after the horizontal rule below."
|
||||
"\n---\n",
|
||||
file_content,
|
||||
]
|
||||
messages.append(LlamaMessage("user", "\n".join(message_lines)))
|
||||
|
||||
new_content = await replicate_llama_completion(app, messages)
|
||||
|
||||
# Remove backtick fences
|
||||
md_code_pattern = re.compile(r"^\s*```[\w\s]*\n([\s\S]*)\n```\s*$", re.MULTILINE)
|
||||
md_match = re.search(md_code_pattern, new_content)
|
||||
if md_match:
|
||||
print("agent used backtick fences")
|
||||
new_content = md_match.group(1)
|
||||
|
||||
return new_content
|
||||
+44
-34
@@ -1,11 +1,13 @@
|
||||
import threading
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import ollama
|
||||
from stream import Stream
|
||||
from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec
|
||||
|
||||
from llama_cpp import Llama
|
||||
from openrouter import OpenRouter
|
||||
|
||||
VALID_ROLES = ['user', 'assistant', 'system', 'tool']
|
||||
CONTEXT_WINDOW = 10_000
|
||||
@@ -130,6 +132,8 @@ async def chat_completion(app,
|
||||
result = await replicate_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
|
||||
elif model_engine == "ollama":
|
||||
result = await ollama_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
|
||||
elif model_engine == "openrouter":
|
||||
result = await openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
|
||||
else:
|
||||
raise Exception("Unsupported model: " + repr(model))
|
||||
|
||||
@@ -236,6 +240,46 @@ async def ollama_chat_completion(app, messages, model, tools=[], opts={}, agent_
|
||||
return content_stream
|
||||
|
||||
|
||||
async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
|
||||
if agent_run:
|
||||
content_stream = agent_run.stream
|
||||
else:
|
||||
content_stream = Stream(app)
|
||||
|
||||
async def worker():
|
||||
print("START OPENROUTER INFERENCE")
|
||||
if not app.openrouter_client:
|
||||
raise Exception("OPENROUTER_API_KEY environment variable is not set")
|
||||
|
||||
kwargs = {
|
||||
'model': model,
|
||||
'messages': [msg.to_dict() for msg in messages],
|
||||
'stream': True,
|
||||
}
|
||||
|
||||
# Pass tools in OpenAI-compatible format if provided
|
||||
formatted_tools = format_tools_for_engine(tools, "openrouter") if tools else None
|
||||
if formatted_tools:
|
||||
kwargs['tools'] = formatted_tools
|
||||
|
||||
stream_response = await app.openrouter_client.chat.send_async(**kwargs)
|
||||
|
||||
async for event in stream_response:
|
||||
if event.choices and len(event.choices) > 0:
|
||||
delta = event.choices[0].delta
|
||||
if delta.content:
|
||||
content_stream.feed(delta.content)
|
||||
|
||||
content_stream.close()
|
||||
print("END OPENROUTER INFERENCE")
|
||||
|
||||
asyncio.create_task(worker())
|
||||
|
||||
if agent_run:
|
||||
return None
|
||||
return content_stream
|
||||
|
||||
|
||||
async def replicate_completion(app, messages, model=None, tools=[], opts={}, agent_run=None):
|
||||
if model is None:
|
||||
model = Models.llama_8b
|
||||
@@ -263,40 +307,6 @@ async def replicate_completion(app, messages, model=None, tools=[], opts={}, age
|
||||
return None
|
||||
return result
|
||||
|
||||
# DEPRECATED
|
||||
async def replicate_llama_completion(app, messages, model=None, opts={}):
|
||||
if model is None:
|
||||
model = Models.llama_8b
|
||||
|
||||
return await replicate_base_completion(app, model, {
|
||||
"prompt": format_instruct(messages),
|
||||
"temperature": opts.get('temperature', 0.75),
|
||||
"max_new_tokens": opts.get('max_tokens', 2048),
|
||||
"system_prompt": "" # System prompt is formatted in the prompt
|
||||
})
|
||||
|
||||
# DEPRECATED
|
||||
async def replicate_qwen_completion(app, messages, model=None, opts={}):
|
||||
if model is None:
|
||||
model = Models.qwen3_8b
|
||||
|
||||
if model == Models.qwen3_8b:
|
||||
input = {
|
||||
"prompt": format_instruct(messages),
|
||||
"temperature": opts.get('temperature', 0.75),
|
||||
"max_new_tokens": opts.get('max_tokens', 2048),
|
||||
}
|
||||
elif model == Models.qwen3_235b:
|
||||
input = {
|
||||
"prompt": format_instruct(messages),
|
||||
"temperature": opts.get('temperature', 0.75),
|
||||
"max_tokens": opts.get('max_tokens', 2048),
|
||||
}
|
||||
else:
|
||||
raise Exception("Invalid qwen model: " + model)
|
||||
|
||||
return await replicate_base_completion(app, model, input)
|
||||
|
||||
async def replicate_base_completion(app, model, input, agent_run=None):
|
||||
if agent_run:
|
||||
content_stream = agent_run.stream
|
||||
|
||||
Reference in new issue
Block a user