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())
 
@@ -172,6 +170,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):

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 = [
@@ -36,4 +39,37 @@ async def gral_message(app, system_prompt, message_history: list, file_map: dict
     messages.extend([LlamaMessage(msg['role'], msg['content']) for msg in message_history])
 
     response = await replicate_llama_completion(app, messages)
-    return response
+    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,15 +67,24 @@ def start_server(app):
 
         # Authenticate before dedicating any resources
         auth_req = json.loads(await websocket.receive())
-        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'}))
-            return
 
-        user_id = app.verify_user_credentials(auth_req['username'], auth_req['password'])
-        if user_id is None:
-            await websocket.send(json.dumps({'error': 'invalid credentials'}))
-            return
+        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'}))
+                return
+
+            user_id = app.verify_user_credentials(auth_req['username'], auth_req['password'])
+            if user_id is None:
+                await websocket.send(json.dumps({'error': 'invalid credentials'}))
+                return
 
         send_queue = asyncio.Queue()
         async def sender_task():