Add endpoint for driver

app.py

@@ -33,34 +33,77 @@ def fetch_prompts_from_google_doc():
 
     return prompts
 
-class Prompt(BaseModel):
+class NavigatePrompt(BaseModel):
     message: str
     code: str
 
+class DrivePrompt(BaseModel):
+    instruction: str
+    code: str
+
+@app.post("/drive")
+async def drive(prompt: DrivePrompt):
+    prompts = fetch_prompts_from_google_doc()
+    print("Received POST to /drive")
+    print("Instruction:", prompt.instruction)
+
+    system_prompt = f"""
+    ### Unit Information ###
+    {prompts.get('UNIT_INFORMATION_DRIVER', '')}
+    
+    ### Role Description ###
+    {prompts.get('ROLE_DESCRIPTION_DRIVER', '')}
+    
+    ### Behavioral Instructions ###
+    {prompts.get('BEHAVIORAL_INSTRUCTIONS_DRIVER', '')}
+    """
+
+    user_prompt = f"""
+    ### Instruction ###
+    {prompt.instruction}
+    
+    ### Existing Code ###
+    {prompt.code}
+    """
+
+    chat_completion = client.chat.completions.create(
+        model="meta-llama/Meta-Llama-3-70B-Instruct",
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ],
+        temperature=0.7,
+        max_tokens=1024,
+    )
+
+    reply = chat_completion.choices[0].message.content.strip()
+    print("Drive response:", reply)
+    return {"reply": reply}
+
 @app.post("/navigate")
-async def chat(prompt: Prompt):
+async def chat(prompt: NavigatePrompt):
     prompts = fetch_prompts_from_google_doc()
     print("Received POST request")
     print("Message:", prompt.message)
 
     system_prompt = f"""
     ### Unit Information ###
-    {prompts['UNIT_INFORMATION']}
+    {prompts.get('UNIT_INFORMATION_NAVIGATOR', '')}
     
     ### Role Description ###
-    {prompts['ROLE_DESCRIPTION']}
+    {prompts.get('ROLE_DESCRIPTION_NAVIGATOR', '')}
     
     ### Topic Information ###
-    {prompts['TOPIC_INFORMATION']}
+    {prompts.get('TOPIC_INFORMATION_NAVIGATOR', '')}
     
     ### Task Description ###
-    {prompts['TASK_DESCRIPTION']}
+    {prompts.get('TASK_DESCRIPTION_NAVIGATOR', '')}
     
     ### Reference Solution ###
-    {prompts['REFERENCE_SOLUTION']}
+    {prompts.get('REFERENCE_SOLUTION_NAVIGATOR', '')}
     
     ### Behavioral Instructions ###
-    {prompts['BEHAVIORAL_INSTRUCTIONS']}
+    {prompts.get('BEHAVIORAL_INSTRUCTIONS_NAVIGATOR', '')}
     """
     
     user_prompt = f"""