upgrade backend to support ios

backend/backend.py

@@ -17,7 +17,6 @@ from pygments.lexers import get_lexer_by_name
 from sqlalchemy import JSON, Column, Integer, String, create_engine
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import Session, sessionmaker
-from sqlalchemy.sql import insert, select, text
 from uvicorn import Config, Server
 
 LANGS = [
@@ -39,8 +38,8 @@ class User(Base):
         return f"User(id={self.id}, uid={self.uid}"
 
 
-class History(Base):
-    __tablename__ = "history"
+class AndroidHistory(Base):
+    __tablename__ = "android_history"
 
     id = Column(Integer, primary_key=True, autoincrement=True)
     uid = Column(String, nullable=False)
@@ -48,7 +47,19 @@ class History(Base):
     answer = Column(String, nullable=False)
 
     def __repr__(self):
-        return f"History(id={self.id}, uid={self.uid}, question={self.question}, answer={self.answer}"
+        return f"AndroidHistory(question={self.question}, answer={self.answer}"
+
+
+class BrowserHistory(Base):
+    __tablename__ = "browser_history"
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    machineid = Column(String, nullable=False)
+    uid = Column(String, nullable=False)
+    url = Column(String, nullable=False)
+
+    def __repr__(self):
+        return f"BrowserHistory(machineid={self.machineid}, url={self.url}"
 
 
 # Add a new table to store the commands
@@ -74,12 +85,32 @@ app = FastAPI(debug=True)
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 
 
-# Add a new API endpoint to add commands to the queue
+class RegisterItem(BaseModel):
+    openai_key: str
+
+
 class CommandItem(BaseModel):
     uid: str
     command: str
 
 
+class EventItem(BaseModel):
+    uid: str
+    event: str
+
+
+class AssistItem(BaseModel):
+    uid: str
+    prompt: str
+    version: str
+
+
+class SaveURLItem(BaseModel):
+    uid: str
+    machineid: str
+    url: str
+
+
 @app.post("/add_command")
 async def add_command(item: CommandItem):
     db: Session = SessionLocal()
@@ -90,11 +121,6 @@ async def add_command(item: CommandItem):
     return {"message": "Command added"}
 
 
-class EventItem(BaseModel):
-    uid: str
-    event: str
-
-
 @app.post("/send_event")
 async def send_event(item: EventItem):
     print(f"Received event from {item.uid}:\n{item.event}")
@@ -127,10 +153,6 @@ async def send_event(item: EventItem):
     }
 
 
-class RegisterItem(BaseModel):
-    openai_key: str
-
-
 @app.post("/register")
 async def register(item: RegisterItem):
     db: Session = SessionLocal()
@@ -145,12 +167,6 @@ async def register(item: RegisterItem):
         return {"uid": new_user.uid}
 
 
-class AssistItem(BaseModel):
-    uid: str
-    prompt: str
-    version: str
-
-
 @app.post("/assist")
 async def assist(item: AssistItem):
     db: Session = SessionLocal()
@@ -160,13 +176,15 @@ async def assist(item: AssistItem):
 
     # Call OpenAI
     openai.api_key = user.openai_key
-    response = openai_text_call(item.prompt, model=item.version)
+    response = generate_quick_completion(item.prompt, item.version)
 
     # Update the last time assist was called
     user.last_assist = datetime.now()
 
     # Store the history
-    new_history = History(uid=item.uid, question=item.prompt, answer=response["text"])
+    new_history = AndroidHistory(
+        uid=item.uid, question=item.prompt, answer=json.loads(str(response))
+    )
 
     db.add(new_history)
     db.commit()
@@ -177,27 +195,52 @@ async def assist(item: AssistItem):
 @app.get("/get_history/{uid}")
 async def get_history(uid: str):
     db: Session = SessionLocal()
-    history = db.query(History).filter(History.uid == uid).all()
+    history = db.query(BrowserHistory).filter(BrowserHistory.uid == uid).all()
+    browser_history = db.query(AndroidHistory).filter(AndroidHistory.uid == uid).all()
     commands = db.query(Command).filter(Command.uid == uid).all()
-    if not history:
-        raise HTTPException(status_code=400, detail="No history found for this uid")
-    ret = {"history": [h.__dict__ for h in history]}
-
-    if commands and len(commands) > 0:
-        try:
-            with open(f"{uid}_events.txt", "r") as f:
-                events = f.read().split(",")
-        except FileNotFoundError:
-            events = None
-
-        return {
-            "events": events,
-            "history": [h.__dict__ for h in history],
-            "commands": [c.__dict__ for c in commands],
-        }
 
+    try:
+        with open(f"{uid}_events.txt", "r") as f:
+            events = f.read().split(",")
+    except FileNotFoundError:
+        events = ""
+
+    return {
+        "events": events,
+        "history": [h.__dict__ for h in history],
+        "browser_history": [h.__dict__ for h in browser_history],
+        "commands": [c.__dict__ for c in commands],
+    }
+
+
+@app.post("/saveurl")
+async def saveurl(item: SaveURLItem):
+    db: Session = SessionLocal()
+    new_browser_history = BrowserHistory(
+        uid=item.uid, machineid=item.machineid, url=item.url
+    )
+    db.add(new_browser_history)
+    db.commit()
+    db.refresh(new_browser_history)
+    return {"message": "Browser history saved"}
+
+
+def generate_quick_completion(prompt, model):
+    dropdrown = model.split(":")
+    engine = dropdrown[0]
+    max_tokens = int(dropdrown[1])
+    if "gpt" in model:
+        message = [{"role": "user", "content": prompt}]
+        response = openai.ChatCompletion.create(
+            model=engine,
+            messages=message,
+            temperature=0.2,
+            max_tokens=max_tokens,
+            frequency_penalty=0.0,
+        )
     else:
-        return {"history": [h.__dict__ for h in history]}
+        raise Exception("Unknown model")
+    return response
 
 
 def assist_interface(uid, prompt, gpt_version):