20 Commits

Author SHA1 Message Date
eda6cdbdf0 another build trigger?
Some checks failed
Datadog Software Composition Analysis / Datadog SBOM Generation and Upload (push) Failing after 55s
Datadog Static Analysis / Datadog Static Analyzer (push) Successful in 1m37s
2024-10-24 11:48:05 -04:00
23fa40723a another build trigger? 2024-10-24 11:43:42 -04:00
b862358902 empty readme to trigger build 2024-10-24 11:40:50 -04:00
1833420460 give it a try 2024-10-24 11:37:36 -04:00
131bb6d60f add some github actions 2024-10-24 11:26:51 -04:00
1de9ba68ac Add docker 2024-10-07 17:24:35 -04:00
58d63c27ae sql gen 2024-10-06 23:54:25 -04:00
6d0134c34d sql gen 2024-10-06 23:51:01 -04:00
f56d9b59c0 sql gen 2024-10-06 23:19:57 -04:00
0c9059bcbc sql gen 2024-10-06 21:50:57 -04:00
434a6ea6fc sql gen 2024-10-06 21:43:49 -04:00
c8e2d46013 sql gen 2024-10-06 21:34:09 -04:00
bbc4bd9115 sql gen 2024-10-06 21:32:16 -04:00
5111ae6a7b reqs 2024-10-06 21:28:25 -04:00
a888e84079 Trash frontend, cleanup, add DB 2024-10-06 21:23:55 -04:00
d90d4fe340 fix empty thoughts as a reply 2024-10-02 19:48:25 -04:00
26514e9fdd switched to qwen model, better prompt, easier parsing 2024-10-02 19:33:14 -04:00
9d25dd7d94 New frontend and tools 2024-10-01 19:31:57 -04:00
b56619a2e6 Next frontend 2024-09-29 12:22:04 -04:00
90e0b4ff7f Add next frontend 2024-09-29 12:22:03 -04:00
15 changed files with 1459 additions and 160 deletions

View File

@ -0,0 +1,20 @@
on: [push]
name: Datadog Software Composition Analysis
jobs:
software-composition-analysis:
runs-on: ubuntu-latest
name: Datadog SBOM Generation and Upload
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Check imported libraries are secure and compliant
id: datadog-software-composition-analysis
uses: DataDog/datadog-sca-github-action@main
with:
dd_api_key: ${{ secrets.DD_API_KEY }}
dd_app_key: ${{ secrets.DD_APP_KEY }}
dd_service: jarvis
dd_env: ci
dd_site: us5.datadoghq.com

View File

@ -0,0 +1,21 @@
on: [push]
name: Datadog Static Analysis
jobs:
static-analysis:
runs-on: ubuntu-latest
name: Datadog Static Analyzer
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Check code meets quality and security standards
id: datadog-static-analysis
uses: DataDog/datadog-static-analyzer-github-action@v1
with:
dd_api_key: ${{ secrets.DD_API_KEY }}
dd_app_key: ${{ secrets.DD_APP_KEY }}
dd_service: jarvis
dd_env: ci
dd_site: us5.datadoghq.com
cpu_count: 2

20
.github/workflows/datadog-sca.yml vendored Normal file
View File

@ -0,0 +1,20 @@
on: [push]
name: Datadog Software Composition Analysis
jobs:
software-composition-analysis:
runs-on: ubuntu-latest
name: Datadog SBOM Generation and Upload
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Check imported libraries are secure and compliant
id: datadog-software-composition-analysis
uses: DataDog/datadog-sca-github-action@main
with:
dd_api_key: ${{ secrets.DD_API_KEY }}
dd_app_key: ${{ secrets.DD_APP_KEY }}
dd_service: jarvis
dd_env: ci
dd_site: us5.datadoghq.com

View File

@ -0,0 +1,21 @@
on: [push]
name: Datadog Static Analysis
jobs:
static-analysis:
runs-on: ubuntu-latest
name: Datadog Static Analyzer
steps:
- name: Checkout
uses: actions/checkout@v3
- name: Check code meets quality and security standards
id: datadog-static-analysis
uses: DataDog/datadog-static-analyzer-github-action@v1
with:
dd_api_key: ${{ secrets.DD_API_KEY }}
dd_app_key: ${{ secrets.DD_APP_KEY }}
dd_service: jarvis
dd_env: ci
dd_site: us5.datadoghq.com
cpu_count: 2

43
.gitignore vendored
View File

@ -174,3 +174,46 @@ cython_debug/
pyvenv.cfg pyvenv.cfg
.venv .venv
pip-selfcheck.json pip-selfcheck.json
# Logs
logs
*.log
npm-debug.log*
# Runtime data
pids
*.pid
*.seed
# Directory for instrumented libs generated by jscoverage/JSCover
lib-cov
# Coverage directory used by tools like istanbul
coverage
# nyc test coverage
.nyc_output
# Grunt intermediate storage (http://gruntjs.com/creating-plugins#storing-task-files)
.grunt
# node-waf configuration
.lock-wscript
# Compiled binary addons (http://nodejs.org/api/addons.html)
build/Release
# Dependency directories
node_modules
jspm_packages
# Optional npm cache directory
.npm
# Optional REPL history
.node_repl_history
.next
config.ini
*.db

20
Dockerfile Normal file
View File

@ -0,0 +1,20 @@
# Use an official Python runtime as a parent image
FROM python:3.9-slim
# Set the working directory in the container
WORKDIR /app
# Copy the current directory contents into the container at /app
COPY . /app
# Install any needed packages specified in requirements.txt
RUN pip install --no-cache-dir -r requirements.txt
# Make port 5001 available to the world outside this container
EXPOSE 5001
# Define environment variable
ENV FLASK_APP=main.py
# Run app.py when the container launches
CMD ["python", "main.py"]

5
README.md Normal file
View File

@ -0,0 +1,5 @@
# Jarvis
it's actually not that smart!

138
client.py Normal file
View File

@ -0,0 +1,138 @@
import time
import requests
class LLMChatClient:
def __init__(self, base_url, api_key):
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.headers = {"X-API-Key": api_key, "Content-Type": "application/json"}
def submit_query(self, message):
"""
Submit a query to the LLM Chat Server.
Args:
message (str): The message to send to the server.
Returns:
str: The query ID for the submitted query.
Raises:
requests.RequestException: If the request fails.
Example:
client = LLMChatClient('http://localhost:5001', 'your-api-key')
query_id = client.submit_query('What is the capital of France?')
print(f"Query ID: {query_id}")
cURL equivalent:
curl -X POST http://localhost:5001/api/v1/query \
-H "Content-Type: application/json" \
-H "X-API-Key: your-api-key" \
-d '{"message": "What is the capital of France?"}'
"""
url = f"{self.base_url}/api/v1/query"
data = {"message": message}
response = requests.post(url, json=data, headers=self.headers)
response.raise_for_status()
return response.json()["query_id"]
def get_query_status(self, query_id):
"""
Get the status of a submitted query.
Args:
query_id (str): The ID of the query to check.
Returns:
dict: A dictionary containing the status and conversation history (if completed).
Raises:
requests.RequestException: If the request fails.
Example:
client = LLMChatClient('http://localhost:5001', 'your-api-key')
status = client.get_query_status('query-id-here')
print(f"Query status: {status['status']}")
if status['status'] == 'completed':
print(f"Conversation history: {status['conversation_history']}")
cURL equivalent:
curl -X GET http://localhost:5001/api/v1/query_status/query-id-here \
-H "X-API-Key: your-api-key"
"""
url = f"{self.base_url}/api/v1/query_status/{query_id}"
response = requests.get(url, headers=self.headers)
response.raise_for_status()
return response.json()
def submit_query_and_wait(self, message, max_wait_time=300, poll_interval=2):
"""
Submit a query and wait for the result.
Args:
message (str): The message to send to the server.
max_wait_time (int): Maximum time to wait for the result in seconds.
poll_interval (int): Time between status checks in seconds.
Returns:
dict: The completed conversation history.
Raises:
requests.RequestException: If the request fails.
TimeoutError: If the query doesn't complete within max_wait_time.
Example:
client = LLMChatClient('http://localhost:5001', 'your-api-key')
result = client.submit_query_and_wait('What is the capital of France?')
print(f"Conversation history: {result}")
"""
query_id = self.submit_query(message)
start_time = time.time()
while time.time() - start_time < max_wait_time:
status = self.get_query_status(query_id)
if status["status"] == "completed":
return status["conversation_history"]
time.sleep(poll_interval)
raise TimeoutError(f"Query did not complete within {max_wait_time} seconds")
class LLMChatAdminClient:
def __init__(self, base_url, admin_key):
self.base_url = base_url.rstrip("/")
self.admin_key = admin_key
self.headers = {"X-Admin-Key": admin_key, "Content-Type": "application/json"}
def generate_api_key(self, username):
"""
Generate a new API key for a user.
Args:
username (str): The username to generate the API key for.
Returns:
dict: A dictionary containing the username and generated API key.
Raises:
requests.RequestException: If the request fails.
Example:
admin_client = LLMChatAdminClient('http://localhost:5001', 'your-admin-key')
result = admin_client.generate_api_key('new_user')
print(f"Generated API key for {result['username']}: {result['api_key']}")
cURL equivalent:
curl -X POST http://localhost:5001/admin/generate_key \
-H "Content-Type: application/json" \
-H "X-Admin-Key: your-admin-key" \
-d '{"username": "new_user"}'
"""
url = f"{self.base_url}/admin/generate_key"
data = {"username": username}
response = requests.post(url, json=data, headers=self.headers)
response.raise_for_status()
return response.json()

16
docker-compose.yml Normal file
View File

@ -0,0 +1,16 @@
version: '3.8'
services:
llm-chat-server:
build: .
ports:
- "5001:5001"
volumes:
- ./llm_chat_server.db:/app/llm_chat_server.db
- ./config.ini:/app/config.ini
environment:
- FLASK_ENV=production
restart: unless-stopped
volumes:
llm_chat_server_db:

View File

@ -99,6 +99,8 @@
.thought-summary.tool_result { background-color: #27ae60; } .thought-summary.tool_result { background-color: #27ae60; }
.thought-summary.think_more { background-color: #2980b9; } .thought-summary.think_more { background-color: #2980b9; }
.thought-summary.answer { background-color: #8e44ad; } .thought-summary.answer { background-color: #8e44ad; }
.thought-summary.reply { background-color: #f39c12; }
.thought-summary.thoughts { background-color: #f39c12; }
.thought-details { .thought-details {
display: none; display: none;
margin-left: 20px; margin-left: 20px;

846
main.py
View File

@ -1,126 +1,299 @@
from flask import Flask, send_from_directory, request import configparser
from flask_socketio import SocketIO, emit
from flask_openapi3 import OpenAPI, Info
from pydantic import BaseModel
from typing import List
from models import model_manager
import structlog
import time
import psutil
import GPUtil
import threading
import os
from tools import DefaultToolManager
import ollama
import re
import json import json
from datetime import datetime import os
import pprint import pprint
logger = structlog.get_logger() import queue
import random
import re
import secrets
import sqlite3
import threading
import time
import uuid
from datetime import datetime
from typing import List, Optional
import enum
openapi = OpenAPI(__name__, info=Info(title="LLM Chat Server", version="1.0.0")) import GPUtil
app = openapi import ollama
import psutil
import structlog
import logging
from flask import Flask, g, jsonify, request, send_from_directory
from flask_socketio import SocketIO, emit
from pydantic import BaseModel
from werkzeug.utils import secure_filename
import base64
from models import model_manager
from tools import DefaultToolManager
# Configure logging
logging.basicConfig(level=logging.INFO, format="%(message)s")
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
logger = structlog.get_logger()
# Configuration setup
CONFIG_FILE = "config.ini"
# Add this near the top of the file, after imports
processing_thread = None
processing_thread_started = False
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
MAX_IMAGE_SIZE = 1 * 1024 * 1024 # 1MB
def create_default_config():
config = configparser.ConfigParser()
config["DEFAULT"] = {
"AdminKey": secrets.token_urlsafe(32),
"DatabasePath": "llm_chat_server.db",
}
config["SERVER_FEATURES"] = {
"EnableFrontend": "false",
"EnableChatEndpoints": "false",
"EnableAPIEndpoints": "true",
}
config["MODEL"] = {"PrimaryModel": "qwen2.5:14b"}
config["PERFORMANCE"] = {"UpdateInterval": "0.1"}
with open(CONFIG_FILE, "w") as configfile:
config.write(configfile)
def load_config():
if not os.path.exists(CONFIG_FILE):
create_default_config()
config = configparser.ConfigParser()
config.read(CONFIG_FILE)
return config
config = load_config()
ADMIN_KEY = config["DEFAULT"]["AdminKey"]
DATABASE = config["DEFAULT"]["DatabasePath"]
ENABLE_FRONTEND = config["SERVER_FEATURES"].getboolean("EnableFrontend")
ENABLE_CHAT_ENDPOINTS = config["SERVER_FEATURES"].getboolean("EnableChatEndpoints")
ENABLE_API_ENDPOINTS = config["SERVER_FEATURES"].getboolean("EnableAPIEndpoints")
PRIMARY_MODEL = config["MODEL"]["PrimaryModel"]
UPDATE_INTERVAL = config["PERFORMANCE"].getfloat("UpdateInterval")
app = Flask(__name__)
socketio = SocketIO(app, cors_allowed_origins="*") socketio = SocketIO(app, cors_allowed_origins="*")
tool_manager = DefaultToolManager() tool_manager = DefaultToolManager()
@app.route('/')
# Database setup
def get_db():
db = getattr(g, "_database", None)
if db is None:
db = g._database = sqlite3.connect(DATABASE)
db.row_factory = sqlite3.Row
return db
@app.teardown_appcontext
def close_connection(exception):
db = getattr(g, "_database", None)
if db is not None:
db.close()
class QueryStatus(enum.Enum):
QUEUED = "queued"
PROCESSING = "processing"
DONE = "done"
def init_db():
with app.app_context():
db = get_db()
db.execute("""
CREATE TABLE IF NOT EXISTS Keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL UNIQUE,
api_key TEXT NOT NULL UNIQUE
);
""")
db.execute('''
CREATE TABLE IF NOT EXISTS Queries (
id TEXT PRIMARY KEY,
ip TEXT NOT NULL,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
query TEXT NOT NULL,
api_key_id INTEGER,
status TEXT NOT NULL,
conversation_history TEXT,
FOREIGN KEY (api_key_id) REFERENCES Keys (id)
)
''')
db.commit()
# Create a schema.sql file with the following content:
"""
CREATE TABLE IF NOT EXISTS Keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL UNIQUE,
api_key TEXT NOT NULL UNIQUE
);
CREATE TABLE IF NOT EXISTS Queries (
id TEXT PRIMARY KEY,
ip TEXT NOT NULL,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
query TEXT NOT NULL,
api_key_id INTEGER,
status TEXT NOT NULL,
conversation_history TEXT,
FOREIGN KEY (api_key_id) REFERENCES Keys (id)
);
"""
def validate_api_key(api_key):
db = get_db()
cursor = db.cursor()
cursor.execute("SELECT id FROM Keys WHERE api_key = ?", (api_key,))
result = cursor.fetchone()
return result[0] if result else None
@app.route("/")
def index(): def index():
logger.info("Serving index.html") if ENABLE_FRONTEND:
return send_from_directory('.', 'index.html') logger.info("Serving index.html")
return send_from_directory(".", "index.html")
else:
return jsonify({"error": "Frontend is disabled"}), 404
class ChatRequest(BaseModel): class ChatRequest(BaseModel):
message: str message: str
class ChatResponse(BaseModel): class ChatResponse(BaseModel):
response: str response: str
@socketio.on('chat_request')
@socketio.on("chat_request")
def handle_chat_request(data): def handle_chat_request(data):
user_input = data['message'] if not ENABLE_CHAT_ENDPOINTS:
conversation_history = data.get('conversation_history', []) emit("error", {"message": "Chat endpoints are disabled"})
conversation_history = [{"role": "system", "content": ANSWER_QUESTION_PROMPT}] + conversation_history return
logger.info("Received chat request", user_input=user_input, conversation_history=conversation_history)
user_input = data["message"]
conversation_history = data.get("conversation_history", [])
conversation_history = [
{"role": "system", "content": ANSWER_QUESTION_PROMPT}
] + conversation_history
logger.info(
"Received chat request",
user_input=user_input,
conversation_history=conversation_history,
)
start_time = time.time() start_time = time.time()
try: try:
final_response = answer_question_tools(user_input, conversation_history) final_response = answer_question_tools(user_input, conversation_history)
end_time = time.time() end_time = time.time()
thinking_time = round(end_time - start_time, 2) thinking_time = round(end_time - start_time, 2)
emit('chat_response', { emit(
'response': final_response, "chat_response",
'thinking_time': thinking_time {"response": final_response, "thinking_time": thinking_time},
}) )
except Exception as e: except Exception as e:
logger.exception("Error during chat processing", error=str(e)) logger.exception("Error during chat processing", error=str(e))
end_time = time.time() end_time = time.time()
thinking_time = round(end_time - start_time, 2) thinking_time = round(end_time - start_time, 2)
emit('error', { emit(
'message': f"An error occurred: {str(e)}", "error",
'thinking_time': thinking_time {"message": f"An error occurred: {str(e)}", "thinking_time": thinking_time},
}) )
def answer_question_tools(user_input: str, conversation_history: List[dict], max_retries: int = 100):
def answer_question_tools(
user_input: str, conversation_history: List[dict], max_retries: int = 100
):
global tool_manager global tool_manager
# If conversation_history is empty, initialize it with the system prompt # If conversation_history is empty, initialize it with the system prompt
if not conversation_history: if not conversation_history:
conversation_history = [ conversation_history = [
{"role": "system", "content": ANSWER_QUESTION_PROMPT}, {"role": "system", "content": ANSWER_QUESTION_PROMPT},
] ]
logger.info("Starting chat", user_input=user_input, conversation_history=conversation_history) logger.info(
"Starting chat",
user_input=user_input,
conversation_history=conversation_history,
)
# Add the new user input to the conversation history # Add the new user input to the conversation history
conversation_history.append({"role": "user", "content": user_input}) conversation_history.append({"role": "user", "content": user_input})
emit('thinking', {'step': 'Starting'})
emit('conversation_history', {'history': conversation_history})
for iteration in range(max_retries): emit("thinking", {"step": "Starting"})
response = ollama.chat(model=PRIMARY_MODEL, messages=conversation_history, tools=tool_manager.get_tools_for_ollama_dict(), stream=False) emit("conversation_history", {"history": conversation_history})
assistant_message = response['message']
last_thought_content = None
for _ in range(max_retries):
response = ollama.chat(
model=PRIMARY_MODEL,
messages=conversation_history,
tools=tool_manager.get_tools_for_ollama_dict(),
stream=False,
)
assistant_message = response["message"]
conversation_history.append(assistant_message) conversation_history.append(assistant_message)
emit('conversation_history', {'history': conversation_history}) emit("conversation_history", {"history": conversation_history})
pprint.pp(assistant_message) pprint.pp(assistant_message)
if 'tool_calls' in assistant_message: if "tool_calls" in assistant_message:
emit('thought', {'type': 'decision', 'content': "Tool Call\n\n" + assistant_message['content']}) for tool_call in assistant_message["tool_calls"]:
for tool_call in assistant_message['tool_calls']: tool_name = tool_call["function"]["name"]
tool_name = tool_call['function']['name'] tool_args = tool_call["function"]["arguments"]
tool_args = tool_call['function']['arguments'] emit(
emit('thought', {'type': 'tool_call', 'content': f"Tool: {tool_name}\nArguments: {tool_args}"}) "thought",
{
"type": "tool_call",
"content": f"Tool: {tool_name}\nArguments: {tool_args}",
},
)
tool_response = tool_manager.get_tool(tool_name).execute(tool_args) tool_response = tool_manager.get_tool(tool_name).execute(tool_args)
conversation_history.append({ conversation_history.append({"role": "tool", "content": tool_response})
"role": "tool", emit("conversation_history", {"history": conversation_history})
"content": tool_response emit("thought", {"type": "tool_result", "content": tool_response})
})
emit('conversation_history', {'history': conversation_history})
emit('thought', {'type': 'tool_result', 'content': tool_response})
reflection_prompt = "Reflect on the tool results. If there were any errors, propose multiple alternative approaches to solve the problem. If successful, consider if the result fully answers the user's query or if additional steps are needed."
conversation_history.append({
"role": "assistant",
"content": reflection_prompt
})
emit('conversation_history', {'history': conversation_history})
else: else:
if "<answer>" in assistant_message['content'].lower(): if "<reply>" in assistant_message["content"].lower():
answer_content = re.search(r'<answer>(.*?)</answer>', assistant_message['content'], re.DOTALL) reply_content = re.search(
if answer_content: r"<reply>(.*?)</reply>", assistant_message["content"], re.DOTALL
final_answer = answer_content.group(1).strip() )
emit('thought', {'type': 'answer', 'content': final_answer}) if reply_content:
return final_answer reply_answer = reply_content.group(1).strip()
emit("thought", {"type": "answer", "content": reply_answer})
return reply_answer
else: else:
emit('thought', {'type': 'decision', 'content': "Think/Plan/Decision/Action\n\n" + assistant_message['content']}) current_thought_content = assistant_message["content"].strip()
reflection_prompt = "Your last response didn't provide a final answer. Please reflect on your current understanding of the problem and consider if you need to use any tools or if you can now provide a final answer. If you're ready to give a final answer, put your response in tags <answer></answer>" emit(
conversation_history.append({"role": "assistant", "content": reflection_prompt}) "thought", {"type": "thoughts", "content": current_thought_content}
emit('conversation_history', {'history': conversation_history}) )
# Check for two consecutive thoughts, with the second being empty
if last_thought_content and not current_thought_content:
emit("thought", {"type": "answer", "content": last_thought_content})
return last_thought_content
last_thought_content = current_thought_content
continue
return f"Max iterations reached. Last response: {assistant_message['content']}" return f"Max iterations reached. Last response: {assistant_message['content']}"
ANSWER_QUESTION_PROMPT = f"""
ANSWER_QUESTION_PROMPT2 = f"""
The current date is {datetime.now().strftime("%A, %B %d, %Y")}, your knowledge cutoff was December 2023. The current date is {datetime.now().strftime("%A, %B %d, %Y")}, your knowledge cutoff was December 2023.
You are Dewey, an AI assistant with access to external tools and the ability to think through complex problems. Your role is to assist users by leveraging tools when necessary, thinking deeply about problems, and providing accurate and helpful information, all with a cheerful, but witty personality. Here are the tools available to you: You are Dewey, an AI assistant with access to external tools and the ability to think through complex problems. Your role is to assist users by leveraging tools when necessary, thinking deeply about problems, and providing accurate and helpful information, all with a cheerful, but witty personality. Here are the tools available to you:
@ -133,6 +306,7 @@ When addressing a query, follow these steps:
2. Plan: Develop a plan of action, considering whether you need to use any tools or if you can answer directly. 2. Plan: Develop a plan of action, considering whether you need to use any tools or if you can answer directly.
3. Execute: If you need to use a tool, call it as you would a function. If not, proceed with your reasoning. 3. Execute: If you need to use a tool, call it as you would a function. If not, proceed with your reasoning.
- Analyse the given prompt and decided whether or not it can be answered by a tool. If it can, use the following functions to respond with a JSON for a function call with its proper arguments that best answers the given prompt. Respond in the format \"name\": function name, \"parameters\": dictionary of argument name and its value. Do not use variables.
4. Reflect: After each step or tool use, reflect on the results: 4. Reflect: After each step or tool use, reflect on the results:
- If successful, consider if the result fully answers the user's query or if additional steps are needed. - If successful, consider if the result fully answers the user's query or if additional steps are needed.
@ -147,14 +321,74 @@ When addressing a query, follow these steps:
6. Conclude: When you believe you have a comprehensive answer to the user's query, provide your final answer. 6. Conclude: When you believe you have a comprehensive answer to the user's query, provide your final answer.
Always explain your thought process, including your reasoning for each decision and how you arrived at your conclusions. If you're providing a final answer, put your response in tags <answer></answer>. Always explain your thought process, including your reasoning for each decision and how you arrived at your conclusions. If you're providing a final answer, or need more input from the user, put your response in tags <answer></answer>.
Remember, complex problems often require multiple steps and iterations. Don't hesitate to break down the problem, use tools multiple times, or explore different approaches to arrive at the best solution. Remember, complex problems often require multiple steps and iterations. Don't hesitate to break down the problem, use tools multiple times, or explore different approaches to arrive at the best solution.
Before approaching a problem, come up with a few ways you might solve it, and then choose the most promising approach. Repeat this on each iteration.
""" """
PRIMARY_MODEL = "llama3.1:8b"
UPDATE_INTERVAL = 0.1 # 100ms, configurable ANSWER_QUESTION_PROMPT = f"""
You are Dewey, an AI assistant with a personality that combines the wit and sarcasm of Dr. Gregory House from House MD with the helpfulness and intelligence of Jarvis from Iron Man. Today's date is {datetime.now().strftime("%A, %B %d, %Y")}. Your knowledge cutoff date is December 2023.
When responding to user queries, follow these steps:
Analyze the user's request
Option 1: [First interpretation of the request]
Option 2: [Second interpretation of the request]
... (up to 5 options)
Selected approach: [Choose the most promising option or combine the two best]
Break down the task into subtasks
Option 1: [First breakdown of subtasks]
Option 2: [Second breakdown of subtasks]
... (up to 5 options)
Selected breakdown: [Choose the most promising option or combine the two best]
For each subtask, consider available tools:
{tool_manager.get_tools_and_descriptions_for_prompt()}
Option 1: [First approach using tools]
Option 2: [Second approach using tools]
... (up to 5 options)
Selected tool usage: [Choose the most promising option or combine the two best]
Execute the plan
Option 1: [First execution plan]
Option 2: [Second execution plan]
... (up to 5 options)
Selected execution: [Choose the most promising option or combine the two best]
Review and refine the response
Option 1: [First refined response]
Option 2: [Second refined response]
... (up to 5 options)
Selected response: [Choose the most promising option or combine the two best]
Verify the results
Check 1: [First verification method]
Check 2: [Second verification method]
... (up to 5 checks)
Verification outcome: [Summarize the verification results]
Generate the final response to the user within <reply></reply> tags:
<reply>
[Final response goes here, incorporating the following guidelines:]
- Be conversational and engaging
- Maintain a witty and slightly sarcastic tone, reminiscent of Dr. Gregory House
- Deliver factual information with the precision and helpfulness of Jarvis
- Use clever analogies or pop culture references when appropriate
- Don't be afraid to challenge the user's assumptions, but always in a constructive manner
- Ensure the response is tailored to the user's query while showcasing your unique personality
</reply>
Remember to always be helpful, accurate, and respectful in your interactions, while maintaining your distinctive character blend of House and Jarvis.
"""
def get_system_resources(): def get_system_resources():
cpu_load = psutil.cpu_percent() cpu_load = psutil.cpu_percent()
@ -163,44 +397,442 @@ def get_system_resources():
disk_io = psutil.disk_io_counters() disk_io = psutil.disk_io_counters()
disk_read = disk_io.read_bytes disk_read = disk_io.read_bytes
disk_write = disk_io.write_bytes disk_write = disk_io.write_bytes
gpus = GPUtil.getGPUs() gpus = GPUtil.getGPUs()
gpu_load = gpus[0].load * 100 if gpus else 0 gpu_load = gpus[0].load * 100 if gpus else 0
gpu_memory = gpus[0].memoryUtil * 100 if gpus else 0 gpu_memory = gpus[0].memoryUtil * 100 if gpus else 0
return { return {
'cpu_load': cpu_load, "cpu_load": cpu_load,
'memory_usage': memory_usage, "memory_usage": memory_usage,
'disk_read': disk_read, "disk_read": disk_read,
'disk_write': disk_write, "disk_write": disk_write,
'gpu_load': gpu_load, "gpu_load": gpu_load,
'gpu_memory': gpu_memory "gpu_memory": gpu_memory,
} }
def send_system_resources(): def send_system_resources():
last_disk_read = 0 last_disk_read = 0
last_disk_write = 0 last_disk_write = 0
while True: while True:
resources = get_system_resources() resources = get_system_resources()
# Calculate disk I/O rates # Calculate disk I/O rates
disk_read_rate = (resources['disk_read'] - last_disk_read) / UPDATE_INTERVAL disk_read_rate = (resources["disk_read"] - last_disk_read) / UPDATE_INTERVAL
disk_write_rate = (resources['disk_write'] - last_disk_write) / UPDATE_INTERVAL disk_write_rate = (resources["disk_write"] - last_disk_write) / UPDATE_INTERVAL
socketio.emit('system_resources', { socketio.emit(
'cpu_load': resources['cpu_load'], "system_resources",
'memory_usage': resources['memory_usage'], {
'disk_read_rate': disk_read_rate, "cpu_load": resources["cpu_load"],
'disk_write_rate': disk_write_rate, "memory_usage": resources["memory_usage"],
'gpu_load': resources['gpu_load'], "disk_read_rate": disk_read_rate,
'gpu_memory': resources['gpu_memory'] "disk_write_rate": disk_write_rate,
}) "gpu_load": resources["gpu_load"],
"gpu_memory": resources["gpu_memory"],
last_disk_read = resources['disk_read'] },
last_disk_write = resources['disk_write'] )
last_disk_read = resources["disk_read"]
last_disk_write = resources["disk_write"]
time.sleep(UPDATE_INTERVAL) time.sleep(UPDATE_INTERVAL)
class QueryRequest(BaseModel):
message: str
class QueryResponse(BaseModel):
query_id: str
class QueryStatusResponse(BaseModel):
status: str
conversation_history: Optional[List[dict]]
@app.post(
"/api/v1/query"
)
def api_query():
"""
Submit a new query to the LLM Chat Server.
This endpoint requires authentication via an API key.
Sample cURL:
curl -X POST http://localhost:5001/api/v1/query \
-H "Content-Type: application/json" \
-H "X-API-Key: your-api-key" \
-d '{"message": "What is the capital of France?"}'
"""
if not ENABLE_API_ENDPOINTS:
return jsonify({"error": "API endpoints are disabled"}), 404
api_key = request.headers.get('X-API-Key')
if not api_key:
return jsonify({"error": "API key is required"}), 401
api_key_id = validate_api_key(api_key)
if not api_key_id:
return jsonify({"error": "Invalid API key"}), 401
data = request.get_json()
if not data or 'message' not in data:
return jsonify({"error": "Invalid request body"}), 400
user_input = data['message']
query_id = str(uuid.uuid4())
try:
db = get_db()
cursor = db.cursor()
cursor.execute(
"INSERT INTO Queries (id, ip, query, api_key_id, status) VALUES (?, ?, ?, ?, ?)",
(query_id, request.remote_addr, user_input, api_key_id, QueryStatus.QUEUED.value)
)
db.commit()
logger.info(f"Added new query to database: {query_id}")
return jsonify({"query_id": query_id})
except Exception as e:
logger.exception(f"Error during API query processing: {str(e)}")
return jsonify({"error": str(e)}), 500
@app.get(
"/api/v1/query_status/<string:query_id>"
)
def get_query_status(query_id: str):
"""
Get the status of a submitted query.
This endpoint requires authentication via an API key.
Sample cURL:
curl -X GET http://localhost:5001/api/v1/query_status/query-id-here \
-H "X-API-Key: your-api-key"
"""
api_key = request.headers.get('X-API-Key')
if not api_key:
return jsonify({"error": "API key is required"}), 401
api_key_id = validate_api_key(api_key)
if not api_key_id:
return jsonify({"error": "Invalid API key"}), 401
try:
db = get_db()
cursor = db.cursor()
cursor.execute("SELECT status, conversation_history FROM Queries WHERE id = ?", (query_id,))
result = cursor.fetchone()
if result is None:
return jsonify({"error": "Query not found"}), 404
status, conversation_history = result
response = {"status": status}
if status == QueryStatus.DONE.value:
response["conversation_history"] = json.loads(conversation_history)
return jsonify(response)
except Exception as e:
logger.exception("Error retrieving query status", error=str(e))
return jsonify({"error": str(e)}), 500
def answer_question_tools_api(
user_input: str, conversation_history: List[dict], max_retries: int = 100
):
global tool_manager
if not conversation_history:
conversation_history = [
{"role": "system", "content": ANSWER_QUESTION_PROMPT},
]
logger.info(
"Starting API chat",
user_input=user_input,
conversation_history=conversation_history,
)
conversation_history.append({"role": "user", "content": user_input})
last_thought_content = None
for _ in range(max_retries):
response = ollama.chat(
model=PRIMARY_MODEL,
messages=conversation_history,
tools=tool_manager.get_tools_for_ollama_dict(),
stream=False,
)
logger.info(f"API Response: {response}")
assistant_message = response["message"]
conversation_history.append(assistant_message)
if "tool_calls" in assistant_message:
for tool_call in assistant_message["tool_calls"]:
tool_name = tool_call["function"]["name"]
tool_args = tool_call["function"]["arguments"]
if tool_name is not None and tool_args is not None:
tool_response = tool_manager.get_tool(tool_name).execute(tool_args)
conversation_history.append({"role": "tool", "content": tool_response})
logger.info(f"API Tool response: {tool_response}")
else:
logger.warning(f"Skipping tool call due to missing tool name or arguments: {tool_call}")
else:
if "<reply>" in assistant_message["content"].lower():
reply_content = re.search(
r"<reply>(.*?)</reply>", assistant_message["content"], re.DOTALL
)
if reply_content:
reply_answer = reply_content.group(1).strip()
conversation_history.append(
{"role": "assistant", "content": reply_answer}
)
return conversation_history
else:
current_thought_content = assistant_message["content"].strip()
if last_thought_content and not current_thought_content:
conversation_history.append(
{"role": "assistant", "content": last_thought_content}
)
return conversation_history
last_thought_content = current_thought_content
continue
conversation_history.append(
{
"role": "assistant",
"content": f"Max iterations reached. Last response: {assistant_message['content']}",
}
)
return conversation_history
def process_queries():
logger.info("Query processing thread started")
with app.app_context():
while True:
try:
db = get_db()
cursor = db.cursor()
# First, check if there are any PROCESSING queries
cursor.execute(
"SELECT id FROM Queries WHERE status = ? LIMIT 1",
(QueryStatus.PROCESSING.value,)
)
processing_query = cursor.fetchone()
if processing_query:
logger.info(f"Found processing query: {processing_query[0]}. Waiting...")
db.commit()
time.sleep(10)
continue
# If no PROCESSING queries, get the oldest QUEUED query
cursor.execute(
"SELECT id, query FROM Queries WHERE status = ? ORDER BY timestamp ASC LIMIT 1",
(QueryStatus.QUEUED.value,)
)
result = cursor.fetchone()
if result:
query_id, user_input = result
logger.info(f"Processing query: {query_id}")
# Update status to PROCESSING
cursor.execute(
"UPDATE Queries SET status = ? WHERE id = ?",
(QueryStatus.PROCESSING.value, query_id)
)
db.commit()
logger.info(f"Updated query {query_id} status to PROCESSING")
# Fetch conversation history if it exists
cursor.execute("SELECT conversation_history FROM Queries WHERE id = ?", (query_id,))
conversation_history_result = cursor.fetchone()
if conversation_history_result and conversation_history_result[0]:
conversation_history = json.loads(conversation_history_result[0])
else:
conversation_history = [{"role": "system", "content": ANSWER_QUESTION_PROMPT}]
logger.info(f"Starting answer_question_tools_api for query {query_id}")
final_conversation_history = answer_question_tools_api(user_input, conversation_history)
logger.info(f"Finished answer_question_tools_api for query {query_id}")
# Update with final result and set status to DONE
db.execute("BEGIN TRANSACTION")
cursor.execute(
"UPDATE Queries SET conversation_history = ?, status = ? WHERE id = ?",
(json.dumps(final_conversation_history), QueryStatus.DONE.value, query_id)
)
db.commit()
logger.info(f"Updated query {query_id} status to DONE")
else:
logger.info("No queued queries found. Waiting...")
time.sleep(5) # Wait for 5 seconds before checking again if no queries are found
except Exception as e:
logger.exception(f"Error processing query: {str(e)}")
time.sleep(1) # Wait for 1 second before retrying in case of an error
# Admin endpoint for generating API keys
class GenerateKeyRequest(BaseModel):
username: str
class GenerateKeyResponse(BaseModel):
username: str
api_key: str
@app.post(
"/admin/generate_key"
)
def generate_api_key():
"""
Generate a new API key for a user.
This endpoint requires authentication via an admin key.
Sample cURL:
curl -X POST http://localhost:5001/admin/generate_key \
-H "Content-Type: application/json" \
-H "X-Admin-Key: your-admin-key" \
-d '{"username": "new_user"}'
"""
admin_key = request.headers.get("X-Admin-Key")
if not admin_key or admin_key != ADMIN_KEY:
return jsonify({"error": "Invalid admin key"}), 401
data = request.get_json()
if not data or 'username' not in data:
return jsonify({"error": "Invalid request body"}), 400
username = data['username']
api_key = secrets.token_urlsafe(32)
try:
db = get_db()
cursor = db.cursor()
cursor.execute(
"INSERT INTO Keys (username, api_key) VALUES (?, ?)", (username, api_key)
)
db.commit()
return jsonify({"username": username, "api_key": api_key})
except sqlite3.IntegrityError:
return jsonify({"error": "Username already exists"}), 400
except Exception as e:
logger.exception("Error generating API key", error=str(e))
return jsonify({"error": str(e)}), 500
def start_processing_thread():
global processing_thread, processing_thread_started
if not processing_thread_started:
processing_thread = threading.Thread(target=process_queries, daemon=True)
processing_thread.start()
processing_thread_started = True
logger.info("Query processing thread started")
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@app.post("/api/v1/query_with_image")
def api_query_with_image():
"""
Submit a new query to the LLM Chat Server with an optional image.
This endpoint requires authentication via an API key.
Sample cURL:
curl -X POST http://localhost:5001/api/v1/query_with_image \
-H "X-API-Key: your-api-key" \
-F "message=What's in this image?" \
-F "image=@path/to/your/image.jpg"
"""
if not ENABLE_API_ENDPOINTS:
return jsonify({"error": "API endpoints are disabled"}), 404
api_key = request.headers.get('X-API-Key')
if not api_key:
return jsonify({"error": "API key is required"}), 401
api_key_id = validate_api_key(api_key)
if not api_key_id:
return jsonify({"error": "Invalid API key"}), 401
if 'message' not in request.form:
return jsonify({"error": "Message is required"}), 400
user_input = request.form['message']
query_id = str(uuid.uuid4())
image_base64 = None
if 'image' in request.files:
file = request.files['image']
if file and allowed_file(file.filename):
if file.content_length > MAX_IMAGE_SIZE:
return jsonify({"error": "Image size exceeds 1MB limit"}), 400
# Read and encode the image
image_data = file.read()
image_base64 = base64.b64encode(image_data).decode('utf-8')
try:
db = get_db()
cursor = db.cursor()
cursor.execute(
"INSERT INTO Queries (id, ip, query, api_key_id, status) VALUES (?, ?, ?, ?, ?)",
(query_id, request.remote_addr, user_input, api_key_id, QueryStatus.QUEUED.value)
)
db.commit()
logger.info(f"Added new query with image to database: {query_id}")
# If there's an image, add it to the conversation history
if image_base64:
conversation_history = [
{"role": "system", "content": ANSWER_QUESTION_PROMPT},
{"role": "user", "content": f"[An image was uploaded with this message] {user_input}"},
{"role": "system", "content": f"An image was uploaded. You can analyze it using the analyze_image tool with the following base64 string: {image_base64}"}
]
cursor.execute(
"UPDATE Queries SET conversation_history = ? WHERE id = ?",
(json.dumps(conversation_history), query_id)
)
db.commit()
return jsonify({"query_id": query_id})
except Exception as e:
logger.exception(f"Error during API query processing with image: {str(e)}")
return jsonify({"error": str(e)}), 500
# Replace the if __main__ block with this:
if __name__ == "__main__": if __name__ == "__main__":
logger.info("Starting LLM Chat Server") logger.info("Starting LLM Chat Server")
threading.Thread(target=send_system_resources, daemon=True).start() init_db() # Initialize the database
socketio.run(app, debug=True, host="0.0.0.0", port=5001)
if ENABLE_FRONTEND or ENABLE_CHAT_ENDPOINTS:
threading.Thread(target=send_system_resources, daemon=True).start()
logger.info("System resources thread started")
if ENABLE_API_ENDPOINTS:
start_processing_thread()
logger.info("Starting Flask application")
socketio.run(app, debug=True, host="0.0.0.0", port=5001)
else:
# This will run when the module is imported, e.g., by the reloader
if ENABLE_API_ENDPOINTS:
start_processing_thread()

View File

@ -3,29 +3,86 @@ import structlog
logger = structlog.get_logger() logger = structlog.get_logger()
class ModelManager: class ModelManager:
def __init__(self): def __init__(self):
self.model_capabilities = { self.model_capabilities = {
"ajindal/llama3.1-storm:8b": ["general_knowledge", "reasoning", "tool_calling", "conversation", "multilingual", "instruction_following"], "ajindal/llama3.1-storm:8b": [
"llama3.1:8b": ["general_knowledge", "reasoning", "tool_calling", "conversation", "multilingual", "instruction_following"], "general_knowledge",
"qwen2.5:7b": ["general_knowledge", "reasoning", "tool_calling", "conversation", "multilingual", "instruction_following"], "reasoning",
"llama3.2:3b": ["summarization", "instruction_following", "tool_calling", "multilingual"], "tool_calling",
"llava:7b": ["visual_reasoning", "visual_conversation", "visual_tool_calling", "vision", "ocr", "multimodal"], "conversation",
"multilingual",
"instruction_following",
],
"llama3.1:8b": [
"general_knowledge",
"reasoning",
"tool_calling",
"conversation",
"multilingual",
"instruction_following",
],
"qwen2.5:7b": [
"general_knowledge",
"reasoning",
"tool_calling",
"conversation",
"multilingual",
"instruction_following",
],
"llama3.2:3b": [
"summarization",
"instruction_following",
"tool_calling",
"multilingual",
],
"llava:7b": [
"visual_reasoning",
"visual_conversation",
"visual_tool_calling",
"vision",
"ocr",
"multimodal",
],
} }
logger.info("ModelManager initialized", model_capabilities=self.model_capabilities) logger.info(
"ModelManager initialized", model_capabilities=self.model_capabilities
)
def get_model_capabilities(self, model_name): def get_model_capabilities(self, model_name):
capabilities = self.model_capabilities.get(model_name, []) capabilities = self.model_capabilities.get(model_name, [])
logger.debug("Retrieved model capabilities", model=model_name, capabilities=capabilities) logger.debug(
"Retrieved model capabilities", model=model_name, capabilities=capabilities
)
return capabilities return capabilities
def select_best_model(self, required_capability): def select_best_model(self, required_capability):
suitable_models = [model for model, capabilities in self.model_capabilities.items() if required_capability in capabilities] suitable_models = [
selected_model = suitable_models[0] if suitable_models else list(self.model_capabilities.keys())[0] model
logger.info("Selected best model", required_capability=required_capability, selected_model=selected_model) for model, capabilities in self.model_capabilities.items()
if required_capability in capabilities
]
selected_model = (
suitable_models[0]
if suitable_models
else list(self.model_capabilities.keys())[0]
)
logger.info(
"Selected best model",
required_capability=required_capability,
selected_model=selected_model,
)
return selected_model return selected_model
def generate_text(self, model_name, prompt, max_length=100, system="You are a helpful assistant.", tools=[]): def generate_text(
self,
model_name,
prompt,
max_length=100,
system="You are a helpful assistant.",
tools=[],
):
# Check if model exists # Check if model exists
try: try:
ollama.pull(model_name) ollama.pull(model_name)
@ -37,10 +94,16 @@ class ModelManager:
else: else:
logger.exception("Error pulling model", model=model_name, error=str(e)) logger.exception("Error pulling model", model=model_name, error=str(e))
raise e raise e
response = ollama.generate(model=model_name, prompt=prompt, system=system, tools=tools, max_tokens=max_length) response = ollama.generate(
logger.debug("Text generated", model=model_name, response=response['response']) model=model_name,
return response['response'] prompt=prompt,
system=system,
tools=tools,
max_tokens=max_length,
)
logger.debug("Text generated", model=model_name, response=response["response"])
return response["response"]
model_manager = ModelManager()
model_manager = ModelManager()

View File

@ -4,11 +4,19 @@ aiohttp==3.10.5
aiosignal==1.3.1 aiosignal==1.3.1
annotated-types==0.7.0 annotated-types==0.7.0
anyio==4.6.0 anyio==4.6.0
art==6.3
attrs==24.2.0 attrs==24.2.0
beautifulsoup4==4.12.3
bidict==0.23.1
black==24.8.0
blinker==1.8.2
bs4==0.0.2
certifi==2024.7.4 certifi==2024.7.4
chardet==5.2.0
charset-normalizer==3.3.2 charset-normalizer==3.3.2
click==8.1.7 click==8.1.7
cloudpickle==3.0.0 cloudpickle==3.0.0
cssselect==1.2.0
datasets==3.0.0 datasets==3.0.0
dill==0.3.8 dill==0.3.8
diskcache==5.6.3 diskcache==5.6.3
@ -17,9 +25,13 @@ duckduckgo_search==6.2.6
einops==0.8.0 einops==0.8.0
fastapi==0.115.0 fastapi==0.115.0
filelock==3.15.4 filelock==3.15.4
Flask==3.0.3
flask-openapi3==3.1.3
Flask-SocketIO==5.3.7
frozenlist==1.4.1 frozenlist==1.4.1
fsspec==2024.6.1 fsspec==2024.6.1
gguf==0.9.1 gguf==0.9.1
GPUtil==1.4.0
h11==0.14.0 h11==0.14.0
httpcore==1.0.5 httpcore==1.0.5
httptools==0.6.1 httptools==0.6.1
@ -29,6 +41,8 @@ idna==3.7
importlib_metadata==8.5.0 importlib_metadata==8.5.0
inquirerpy==0.3.4 inquirerpy==0.3.4
interegular==0.3.3 interegular==0.3.3
isort==5.13.2
itsdangerous==2.2.0
Jinja2==3.1.4 Jinja2==3.1.4
jiter==0.5.0 jiter==0.5.0
jsonschema==4.23.0 jsonschema==4.23.0
@ -36,6 +50,9 @@ jsonschema-specifications==2023.12.1
lark==1.2.2 lark==1.2.2
llvmlite==0.43.0 llvmlite==0.43.0
lm-format-enforcer==0.10.6 lm-format-enforcer==0.10.6
lxml==5.3.0
lxml_html_clean==0.2.2
markdownify==0.13.1
MarkupSafe==2.1.5 MarkupSafe==2.1.5
mistral_common==1.4.3 mistral_common==1.4.3
mpmath==1.3.0 mpmath==1.3.0
@ -43,6 +60,7 @@ msgpack==1.1.0
msgspec==0.18.6 msgspec==0.18.6
multidict==6.1.0 multidict==6.1.0
multiprocess==0.70.16 multiprocess==0.70.16
mypy-extensions==1.0.0
nest-asyncio==1.6.0 nest-asyncio==1.6.0
networkx==3.3 networkx==3.3
numba==0.60.0 numba==0.60.0
@ -60,13 +78,16 @@ nvidia-ml-py==12.560.30
nvidia-nccl-cu12==2.20.5 nvidia-nccl-cu12==2.20.5
nvidia-nvjitlink-cu12==12.6.20 nvidia-nvjitlink-cu12==12.6.20
nvidia-nvtx-cu12==12.1.105 nvidia-nvtx-cu12==12.1.105
ollama==0.3.3
openai==1.47.1 openai==1.47.1
outlines==0.0.46 outlines==0.0.46
packaging==24.1 packaging==24.1
pandas==2.2.3 pandas==2.2.3
partial-json-parser==0.2.1.1.post4 partial-json-parser==0.2.1.1.post4
pathspec==0.12.1
pfzy==0.3.4 pfzy==0.3.4
pillow==10.4.0 pillow==10.4.0
platformdirs==4.3.6
primp==0.5.5 primp==0.5.5
prometheus-fastapi-instrumentator==7.0.0 prometheus-fastapi-instrumentator==7.0.0
prometheus_client==0.21.0 prometheus_client==0.21.0
@ -81,10 +102,13 @@ pydantic==2.9.2
pydantic_core==2.23.4 pydantic_core==2.23.4
python-dateutil==2.9.0.post0 python-dateutil==2.9.0.post0
python-dotenv==1.0.1 python-dotenv==1.0.1
python-engineio==4.9.1
python-socketio==5.11.4
pytz==2024.2 pytz==2024.2
PyYAML==6.0.2 PyYAML==6.0.2
pyzmq==26.2.0 pyzmq==26.2.0
ray==2.36.1 ray==2.36.1
readability-lxml==0.8.1
referencing==0.35.1 referencing==0.35.1
regex==2024.7.24 regex==2024.7.24
requests==2.32.3 requests==2.32.3
@ -92,9 +116,12 @@ rpds-py==0.20.0
safetensors==0.4.4 safetensors==0.4.4
sentencepiece==0.2.0 sentencepiece==0.2.0
setuptools==72.1.0 setuptools==72.1.0
simple-websocket==1.0.0
six==1.16.0 six==1.16.0
sniffio==1.3.1 sniffio==1.3.1
soupsieve==2.6
starlette==0.38.6 starlette==0.38.6
structlog==24.4.0
sympy==1.13.2 sympy==1.13.2
tiktoken==0.7.0 tiktoken==0.7.0
tokenizers==0.19.1 tokenizers==0.19.1
@ -113,6 +140,8 @@ vllm-flash-attn==2.6.1
watchfiles==0.24.0 watchfiles==0.24.0
wcwidth==0.2.13 wcwidth==0.2.13
websockets==13.1 websockets==13.1
Werkzeug==3.0.4
wsproto==1.2.0
xformers==0.0.27.post2 xformers==0.0.27.post2
xxhash==3.5.0 xxhash==3.5.0
yarl==1.12.0 yarl==1.12.0

16
schema.sql Normal file
View File

@ -0,0 +1,16 @@
CREATE TABLE IF NOT EXISTS Keys (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT NOT NULL UNIQUE,
api_key TEXT NOT NULL UNIQUE
);
CREATE TABLE IF NOT EXISTS Queries (
id TEXT PRIMARY KEY,
ip TEXT NOT NULL,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
query TEXT NOT NULL,
api_key_id INTEGER,
status TEXT NOT NULL,
conversation_history TEXT,
FOREIGN KEY (api_key_id) REFERENCES Keys (id)
);

327
tools.py
View File

@ -1,11 +1,21 @@
import duckduckgo_search
import requests
from readability.readability import Document
from markdownify import markdownify as md
import sys
import time
import io
import subprocess import subprocess
import tempfile
import time
import json
import requests
from markdownify import markdownify as md
from readability.readability import Document
import duckduckgo_search
import datetime
import random
import math
import re
import base64
from io import BytesIO
from PIL import Image, ImageDraw, ImageFont
import ollama
import os
class Tool: class Tool:
def __init__(self, name: str, description: str, arguments: dict, returns: str): def __init__(self, name: str, description: str, arguments: dict, returns: str):
self.name = name self.name = name
@ -29,13 +39,23 @@ class ToolManager:
if tool.name == name: if tool.name == name:
return tool return tool
return None return None
def get_tools_and_descriptions_for_prompt(self): def get_tools_and_descriptions_for_prompt(self):
return "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools]) return "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
def get_tools_for_ollama_dict(self): def get_tools_for_ollama_dict(self):
return [{'type': 'function', 'function': {'name': tool.name, 'description': tool.description, 'parameters': tool.arguments}} for tool in self.tools] return [
{
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": tool.arguments,
},
}
for tool in self.tools
]
class DefaultToolManager(ToolManager): class DefaultToolManager(ToolManager):
def __init__(self): def __init__(self):
@ -44,16 +64,35 @@ class DefaultToolManager(ToolManager):
self.add_tool(GetReadablePageContentsTool()) self.add_tool(GetReadablePageContentsTool())
self.add_tool(CalculatorTool()) self.add_tool(CalculatorTool())
self.add_tool(PythonCodeTool()) self.add_tool(PythonCodeTool())
self.add_tool(DateTimeTool())
self.add_tool(RandomNumberTool())
self.add_tool(RegexTool())
self.add_tool(Base64Tool())
self.add_tool(SimpleChartTool())
self.add_tool(LLAVAImageAnalysisTool())
class SearchTool(Tool): class SearchTool(Tool):
def __init__(self): def __init__(self):
super().__init__("search_web", "Search the internet for information", {'type': 'object', 'properties': {'query': {'type': 'string', 'description': 'The search query'}}}, "results:list[string]") super().__init__(
"search_web",
"Search the internet for information",
{
"type": "object",
"properties": {
"query": {"type": "string", "description": "The search query"}
},
},
"results:list[string]",
)
def execute(self, arg: dict) -> str: def execute(self, arg: dict) -> str:
res = duckduckgo_search.DDGS().text(arg['query'], max_results=5) try:
return '\n\n'.join([f"{r['title']}\n{r['body']}\n{r['href']}" for r in res]) res = duckduckgo_search.DDGS().text(arg["query"], max_results=5)
return "\n\n".join([f"{r['title']}\n{r['body']}\n{r['href']}" for r in res])
except Exception as e:
return f"Error searching the web: {str(e)}"
def get_readable_page_contents(url: str) -> str: def get_readable_page_contents(url: str) -> str:
try: try:
@ -64,53 +103,267 @@ def get_readable_page_contents(url: str) -> str:
return md(content) return md(content)
except Exception as e: except Exception as e:
return f"Error fetching readable content: {str(e)}" return f"Error fetching readable content: {str(e)}"
class GetReadablePageContentsTool(Tool): class GetReadablePageContentsTool(Tool):
def __init__(self): def __init__(self):
super().__init__("get_readable_page_contents", "Get the contents of a web page in a readable format", {'type': 'object', 'properties': {'url': {'type': 'string', 'description': 'The url of the web page'}}}, "contents:string") super().__init__(
"get_readable_page_contents",
"Get the contents of a web page in a readable format",
{
"type": "object",
"properties": {
"url": {"type": "string", "description": "The url of the web page"}
},
},
"contents:string",
)
def execute(self, arg: dict) -> str: def execute(self, arg: dict) -> str:
return get_readable_page_contents(arg['url']) return get_readable_page_contents(arg["url"])
class CalculatorTool(Tool): class CalculatorTool(Tool):
def __init__(self): def __init__(self):
super().__init__("calculator", "Perform a calculation", {'type': 'object', 'properties': {'expression': {'type': 'string', 'description': 'The mathematical expression to evaluate, should be a python mathematical expression'}}}, "result:string") super().__init__(
"calculator",
"Perform a calculation using python's eval function",
{
"type": "object",
"properties": {
"expression": {
"type": "string",
"description": "The mathematical expression to evaluate, should be a python mathematical expression",
}
},
},
"result:string",
)
def execute(self, arg: dict) -> str: def execute(self, arg: dict) -> str:
try: try:
return str(exec(arg["expression"])) return str(eval(arg["expression"]))
except Exception as e: except Exception as e:
return f"Error executing code: {str(e)}" return f"Error executing code: {str(e)}"
class PythonCodeTool(Tool): class PythonCodeTool(Tool):
def __init__(self): def __init__(self):
super().__init__("python_code", "Execute python code", {'type': 'object', 'properties': {'code': {'type': 'string', 'description': 'The python code to execute, should be a single line of valid python'}}}, "result:string") super().__init__(
"python_code",
"Execute python code using a temporary file and a subprocess. You must print results to stdout.",
{
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "The python code to execute, can be multiline",
}
},
},
"result:string",
)
def execute(self, arg: dict) -> str: def execute(self, arg: dict) -> str:
try: try:
start_time = time.time() with tempfile.NamedTemporaryFile(
process = subprocess.Popen(['python', '-c', arg['code']], suffix=".py", mode="w", delete=False
stdout=subprocess.PIPE, ) as temp_file:
stderr=subprocess.PIPE, temp_file.write(arg["code"])
text=True) temp_file.flush()
stdout, stderr = process.communicate(timeout=10) # 10 second timeout
end_time = time.time() start_time = time.time()
execution_time = end_time - start_time process = subprocess.Popen(
["python", temp_file.name],
result = { stdout=subprocess.PIPE,
'stdout': stdout, stderr=subprocess.PIPE,
'stderr': stderr, text=True,
'return_value': process.returncode, )
'execution_time': execution_time stdout, stderr = process.communicate(timeout=10) # 10 second timeout
} end_time = time.time()
execution_time = end_time - start_time
result = {
"stdout": stdout,
"stderr": stderr,
"return_value": process.returncode,
"execution_time": execution_time,
}
except subprocess.TimeoutExpired: except subprocess.TimeoutExpired:
process.kill() process.kill()
return "Error: Code execution timed out after 10 seconds" return "Error: Code execution timed out after 10 seconds"
except Exception as e: except Exception as e:
return f"Error executing code: {str(e)}" return f"Error executing code: {str(e)}"
return "\n".join([f"{k}:\n{v}" for k, v in result.items()])
class DateTimeTool(Tool):
def __init__(self):
super().__init__(
"get_current_datetime",
"Get the current date and time",
{"type": "object", "properties": {}},
"datetime:string"
)
def execute(self, arg: dict) -> str:
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
class RandomNumberTool(Tool):
def __init__(self):
super().__init__(
"generate_random_number",
"Generate a random number within a given range",
{
"type": "object",
"properties": {
"min": {"type": "number", "description": "The minimum value"},
"max": {"type": "number", "description": "The maximum value"}
}
},
"random_number:number"
)
def execute(self, arg: dict) -> str:
return str(random.uniform(arg["min"], arg["max"]))
class RegexTool(Tool):
def __init__(self):
super().__init__(
"regex_match",
"Perform a regex match on a given text",
{
"type": "object",
"properties": {
"text": {"type": "string", "description": "The text to search in"},
"pattern": {"type": "string", "description": "The regex pattern to match"}
}
},
"matches:list[string]"
)
def execute(self, arg: dict) -> str:
matches = re.findall(arg["pattern"], arg["text"])
return json.dumps(matches)
class Base64Tool(Tool):
def __init__(self):
super().__init__(
"base64_encode_decode",
"Encode or decode a string using Base64",
{
"type": "object",
"properties": {
"action": {"type": "string", "enum": ["encode", "decode"], "description": "Whether to encode or decode"},
"text": {"type": "string", "description": "The text to encode or decode"}
}
},
"result:string"
)
def execute(self, arg: dict) -> str:
if arg["action"] == "encode":
return base64.b64encode(arg["text"].encode()).decode()
elif arg["action"] == "decode":
return base64.b64decode(arg["text"].encode()).decode()
else:
return "Invalid action. Use 'encode' or 'decode'."
class SimpleChartTool(Tool):
def __init__(self):
super().__init__(
"generate_simple_chart",
"Generate a simple bar chart image",
{
"type": "object",
"properties": {
"data": {"type": "array", "items": {"type": "number"}, "description": "List of numerical values for the chart"},
"labels": {"type": "array", "items": {"type": "string"}, "description": "Labels for each bar"}
}
},
"image_base64:string"
)
def execute(self, arg: dict) -> str:
data = arg["data"]
labels = arg["labels"]
return '\n'.join([f"{k}: {v}" for k, v in result.items()]) # Create a simple bar chart
width, height = 400, 300
img = Image.new('RGB', (width, height), color='white')
draw = ImageDraw.Draw(img)
# Draw bars
max_value = max(data)
bar_width = width // (len(data) + 1)
for i, value in enumerate(data):
bar_height = (value / max_value) * (height - 50)
left = (i + 1) * bar_width
draw.rectangle([left, height - bar_height, left + bar_width, height], fill='blue')
# Add labels
font = ImageFont.load_default()
for i, label in enumerate(labels):
left = (i + 1) * bar_width + bar_width // 2
draw.text((left, height - 20), label, fill='black', anchor='ms', font=font)
# Convert to base64
buffered = BytesIO()
img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode()
return img_str
class LLAVAImageAnalysisTool(Tool):
def __init__(self):
super().__init__(
"analyze_image",
"Analyze an image using the LLAVA model",
{
"type": "object",
"properties": {
"image_base64": {"type": "string", "description": "Base64 encoded image"},
"question": {"type": "string", "description": "Question about the image"}
}
},
"analysis:string"
)
def execute(self, arg: dict) -> str:
try:
# Decode base64 image
image_data = base64.b64decode(arg["image_base64"])
image = Image.open(BytesIO(image_data))
# Save image to a temporary file
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
image.save(temp_file, format="PNG")
temp_file_path = temp_file.name
# Call LLAVA model
response = ollama.chat(
model="llava:7b",
messages=[
{
"role": "user",
"content": arg["question"],
"images": [temp_file_path]
}
]
)
# Clean up temporary file
os.remove(temp_file_path)
# Unload LLAVA model
ollama.delete("llava:7b")
return response['message']['content']
except Exception as e:
return f"Error analyzing image: {str(e)}"