Token based authentication

backend/agent.py

@@ -56,8 +56,6 @@ class Agent:
         ])
         self.toolset_map = {}
 
-        print("tool prompt section", self.tool_prompt_section)
-
         self.add_toolset(SpecialToolset())
         self.add_toolset(BasicToolset())
 
@@ -173,6 +171,9 @@ class Agent:
     async def gral_message(self, message_history: list, file_map: dict):
         return await gral.gral_message(self.app, self.config['system_prompt'], message_history, file_map)
     
+    async def gral_update_file(self, message_history: list, filepath: str, file_content: str):
+        return await gral.gral_update_file(self.app, self.config['system_prompt'], message_history, filepath, file_content)
+
 
 class IntroAgent(Agent):
     def __init__(self, app, agent_id, _data):

backend/gral.py

@@ -1,4 +1,5 @@
 from inference import LlamaMessage, replicate_llama_completion
+import re
 
 message_format_section = """
 When suggesting code to the user, keep the code as small snippets.
@@ -8,7 +9,9 @@ updated segment
 ```"
 Omit the file link when the block is not related to a known file
 Code blocks must be under 30 lines. Assistant 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."""
+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 = [
@@ -37,3 +40,36 @@ async def gral_message(app, system_prompt, message_history: list, file_map: dict
 
     response = await replicate_llama_completion(app, messages)
     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 raw file contents of the file {filepath} are included following 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

backend/requirements.txt

@@ -1,4 +1,6 @@
 Quart==0.20.0
+quart-auth
+quart-cors
 sqlalchemy
 nest_asyncio
 replicate

backend/server.py

@@ -1,10 +1,23 @@
 from quart import Quart, Response, request, send_from_directory, websocket
+from quart_auth import (
+    QuartAuth, current_user,
+    AuthUser, login_user, logout_user
+)
 import json
 import asyncio
 import nest_asyncio
 
 def start_server(app):
     server = Quart(__name__, static_folder='../frontend/build/web')
+    server.secret_key = 'horde-secret'
+    auth_manager = QuartAuth(server, mode="brearer")
+
+    @server.after_request
+    def after_request(response):
+        response.headers.add('Access-Control-Allow-Origin', '*')
+        response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
+        response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
+        return response
 
     @server.route('/')
     async def index():
@@ -14,6 +27,21 @@ def start_server(app):
     async def serve_static(path):
         return await send_from_directory(server.static_folder, path)
 
+    @server.route('/user/login', methods=['POST'])
+    async def login():
+        data = await request.get_json()
+
+        if "username" not in data or "password" not in data:
+            await websocket.send(json.dumps({'error': 'invalid arguments'}))
+            return
+
+        user_id = app.verify_user_credentials(data['username'], data['password'])
+        if user_id == None:
+            return jsonify({'error': "bad credentials"}), 401
+
+        token = auth_manager.dump_token(user_id)
+        return {"token": token}
+
     @server.route('/user/signup', methods=['POST'])
     async def signup():
         data = await request.get_json()
@@ -39,6 +67,15 @@ def start_server(app):
 
         # Authenticate before dedicating any resources
         auth_req = json.loads(await websocket.receive())
+
+        if 'token' in auth_req:
+            try:
+                user_id = auth_manager.load_token(auth_req['token'])
+            except Exception as err:
+                print("Token authentication failed: ", err)
+                await websocket.send(json.dumps({'error': 'token authentication failed'}))
+                return
+        else:
             if "username" not in auth_req or "password" not in auth_req:
                 print("Invalid auth request:", auth_req)
                 await websocket.send(json.dumps({'error': 'invalid auth'}))

scratchpad.md

@@ -93,63 +93,3 @@ User is only directed to the editing page after its created, has the chance to c
 
 Packages useful for the chat editor:
 - https://pub.dev/packages/custom_text
-
-
-# Gralliot
-
---- Gralliot System Prompt ---
-Your name is Gralliot. You are a Software Engineering AI Assistant. Help the user with any task he asks.
-More than just an AI Assistant, you are a companion to the user in his quests. When its determined that there's missing information, you use your skills to fill the gaps, and when something needs to be done, you proactively use problem solving to determine the best way to use your tools, and immediately executing your solutions.
-You have explicit permission to use your tools at any time, but if there are deeper consequences to your actions, you communicate them to the user.
-
-In the project you're working on, the 'backend' directory contains a Quart python web server, and the 'frontend' directory contains a flutter application. See 'readme.md' for more information.
-
-
-
---- GralMsg Script ---
-
-history.insert(0, {
-  "role":"system",
-  "content":system_prompt,
-  })
-
-history.insert(1, {
-  "role": "user",
-  "content": "# Format Instructions
-When suggesting code to the user, keep the code as small snippets.
-If the update is meant to be applied to a file, include a line linking the file 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. Do not use 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."
-})
-
-file_parts = []
-for filename, content in files.items():
-  file_parts.append(
-    f"
----
-[{filename}]
-
-{content}
-")
-
-# Don't add a file message if there are no files
-if file_parts:
-  preamble = '# Files
-
-Each file section in this message starts with a horizontal rule, followed by the file name in square brackets, followed by the file content.
-'
-
-  history.append({
-    "role": "user",
-    "content": preamble + "
-".join(file_parts),
-    })
-
-history.append({
-  "role": "user",
-  "content": message,
-  })

testing/login.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+URL="http://localhost:5000/user/login"
+DATA='{"username":"test-user","password":"123456"}'
+
+RESPONSE=$(curl -s -X POST $URL -H "Content-Type: application/json" -d "$DATA")
+echo $RESPONSE

testing/suite.js

@@ -2,7 +2,8 @@ const WebSocket = require('ws');
 
 const USE_LLM = false;
 
-const ws = new WebSocket('ws://localhost:5000/ws');
+const hostname = 'localhost:5000';
+const ws = new WebSocket(`ws://${hostname}/ws`);
 var nonceCounter = 1;
 var nonceCallbacks = {};
 var messageQueue = [];
@@ -10,8 +11,27 @@ var messageResolverQueue = [];
 
 // Setup
 
-ws.on('open', function open() {
+const credentials = { username: "test-user", password: "123456" };
-    const packet = JSON.stringify({ username: "test-user", password: "123456" });
+
+const tokenPromise = new Promise(async resolve => {
+    try {
+        const response = await fetch(`http://${hostname}/user/login`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify(credentials),
+        });
+        const data = await response.json();
+        resolve(data.token);
+    } catch (error) {
+        console.error("Login Failed", error);
+        process.exit(1);
+    }
+});
+
+ws.on('open', async function open() {
+    const token = await tokenPromise;
+    console.log("Token:", token);
+    const packet = JSON.stringify({token});
     ws.send(packet);
 
     runTest();
@@ -125,8 +145,9 @@ async function runTest() {
 
     await testChats();
 
-    // Should be enabled with USE_LLM
+    if (USE_LLM) {
         await testGral(mainAgentId);
+    }
 
     process.exit();
 }