dev step aouth....

main.py

@@ -5,16 +5,12 @@ import os, json
 import time
 from openai import OpenAI
 from hashlib import sha256
-from modules import oauth, security, log_module, llm, db
+from modules import model, oauth, security, log_module, llm
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 
 
-
-
-
-
-app = FastAPI()
+app = FastAPI(docs_url=None, redoc_url=None)
 httpsecurity = HTTPBasic()
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
@@ -26,68 +22,47 @@ fecha_unix = str(int(time.time()))
 log_module.log_write("master", "iniciado", "-")
 
 
-
-
-def manejar_funciones(chunk, response_async, data):
-    f_cb = dict(chunk["choices"][0]["delta"]["function_call"])
-    for chunk in response_async:
-        if chunk["choices"][0]["finish_reason"]: 
-            break   
-        for k, v in chunk["choices"][0]["delta"]["function_call"].items():
-            f_cb[k] += v
+@app.get("/", response_class=HTMLResponse)
+async def main_page(request: Request):
     
-    f_cb["arguments"] = json.loads(f_cb["arguments"])
-    log_module.log_write(data["token_data"]["user"], "Funcion ejecutada RQ", json.dumps(f_cb))
-    resultado_funcion = chat_functions.function_callbacks[f_cb["name"]](**f_cb["arguments"])
-    log_module.log_write(data["token_data"]["user"], "Funcion ejecutada RS", json.dumps(resultado_funcion))
-    data["messages"].append({"role": "function",
-                "name": f_cb["name"],
-                "content": resultado_funcion,
-            })
-    return json.dumps(f_cb), llm.ejecutar(data)
+    if not (token := request.cookies.get('token', "")):
+        return RedirectResponse(url='/login')
 
-@app.get("/", response_class=HTMLResponse)
-@app.post("/", response_class=HTMLResponse)
-async def root(request: Request, data = Depends(security.validate_cookie)):
-    data, redirect = data
-    response = None
+    user = model.User.find_from_token(token)
+    
+    if not user.can_use("chat"):
+        return RedirectResponse(url='/hold')
+    
+    token = user.create_token()
+    response = templates.TemplateResponse("main.html", {"request": request, "token": token, "version": fecha_unix})
+    response.set_cookie("token", token)
+    return response
 
-    if redirect:
-        redirect.delete_cookie("token")
-        return redirect
 
-    user = db.check_user(data["token_data"], False)
-    
-    if security.user_scale(user["role"], "chat"):
-        response = templates.TemplateResponse("main.html", {"request": request, "token": data["token"], "version": fecha_unix})
-        response.set_cookie("token", data["token"])
-        
-    elif "on hold" in user["role"]:
-        if request.method=="POST" and not user["description"]:
-            form = await request.form()
-            message = form["message"].strip()
-            
-            db.add_description(user["gid"], message)
-            response = RedirectResponse(url='/')
-        else:
-            response = templates.TemplateResponse("no_access.html", {
-                "request": request,
-                "description": bool(user["description"])
-            })
+########################## SECURITY ##########################
+@app.get("/login", response_class=HTMLResponse)
+async def login(request: Request):
+    ret = templates.TemplateResponse("iniciar_sesion.html", {"request": request, "redirecturi": oauth.OAUTH_REDIRECT})
+    ret.delete_cookie("token")
+    return ret
 
+@app.get("/oauth", response_class=HTMLResponse)
+async def validate_oauth(request: Request):
+    params = dict(request.query_params)
+    google_userinfo = oauth.validate_redirect(params)
+    user = model.User.find_or_create(google_userinfo)
+    token = user.create_token()
+    response = RedirectResponse(url='/')
+    response.set_cookie(key="token", value=token)
     return response
+########################## SECURITY ##########################
 
-        
-
-@app.get("/test", response_class=HTMLResponse)
-async def test(request: Request):
-    return templates.TemplateResponse("no_access.html", {"request": request})
-    
+########################## API ##########################    
 @app.post("/chat")
-async def chat(data = Depends(security.validate_token)):
-    messages = data.get("messages", "")
+async def chat(request: Request):
+    
     
-    if not messages:
+    if not (messages := data.get("messages", "")):
         print("Empty message")
         error = "What ??"
         log_module.log_write(data["token_data"]["user"], "Mensaje RQ", "error, mensajes vacios")
@@ -110,46 +85,52 @@ async def chat(data = Depends(security.validate_token)):
     message = completion.content
     role = completion.role
     yield {"comando": "mensaje", "mensaje": {"role": role, "content": message}}
-
+########################## API ##########################    
     
     
 
-@app.get("/read_log", response_class=HTMLResponse)
-async def read_log(request: Request, credentials: HTTPBasicCredentials = Depends(httpsecurity)):
-    if sha256(credentials.username.encode()).hexdigest()=="bc1c32d709aef061bbde4fc848421cdb933e8a9f391c3a089f2861ac0772c168" and security.authenticate_user(credentials):
-        log_module.log_write(credentials.username, "Log Accesado", "")
-        with open("logs/eventos.log", "r") as f:
-            return HTMLResponse(f.read())
-    log_module.log_write(credentials.username, "Intento acceder logs", f"{request.client.host}:{request.client.port} - {str(dict(request.headers))}")
-    raise HTTPException(status_code=404)    
 
+########################## Static Pages ##########################    
 @app.get("/privacy", response_class=HTMLResponse)
-async def sesion_i(request: Request):
+async def privacy_policy(request: Request):
     return templates.TemplateResponse("PrivacyPolicy.html", {"request": request})
-   
-@app.get("/login", response_class=HTMLResponse)
-async def login(request: Request):
-    ret = templates.TemplateResponse("iniciar_sesion.html", {"request": request, "redirecturi": oauth.OAUTH_REDIRECT})
-    ret.delete_cookie("token")
-    return ret
 
-@app.get("/oauth", response_class=HTMLResponse)
-async def validate_oauth(request: Request):
-    params = dict(request.query_params)
-    google_userinfo = oauth.validate_redirect(params)
+@app.route("/hold", ["GET", "POST"], response_class=HTMLResponse)
+async def on_hold(request: Request):
+    token = request.cookies.get('token', "")
+    redirect_to_login = RedirectResponse(url='/login')
+    if not token:
+        return redirect_to_login
+        
+    
+    user = model.User.find_from_token(token)
 
-    user = db.check_user(google_userinfo)
-    token = security.create_jwt_token({
-        "user": user["name"],
-        "role": user["role"],
-        "gid": user["gid"],
+    if user.can_use("chat"):
+        return redirect_to_login
+    
+    if request.method=="POST" and not user.description:
+        form = await request.form() 
+        if message := form["message"].strip():
+            user.update_description(message)
+    
+    return templates.TemplateResponse(
+        "no_access.html", {
+            "request": request,
+            "description": bool(user.description)
         })
 
-    response = RedirectResponse(url='/')
-    response.set_cookie(key="token", value=token)
-
-    return response
-
+    
+########################## Static Pages ##########################    
 
 
 
+########################## Other ##########################    
+@app.get("/read_log", response_class=HTMLResponse)
+async def read_log(request: Request, credentials: HTTPBasicCredentials = Depends(httpsecurity)):
+    if sha256(credentials.username.encode()).hexdigest()=="bc1c32d709aef061bbde4fc848421cdb933e8a9f391c3a089f2861ac0772c168" and security.authenticate_user(credentials):
+        log_module.log_write(credentials.username, "Log Accesado", "")
+        with open("logs/eventos.log", "r") as f:
+            return HTMLResponse(f.read())
+    log_module.log_write(credentials.username, "Intento acceder logs", f"{request.client.host}:{request.client.port} - {str(dict(request.headers))}")
+    raise HTTPException(status_code=404)    
+########################## Other ##########################    

modules/db.py

@@ -1,41 +0,0 @@
-import os 
-from pymongo.mongo_client import MongoClient
-from pymongo.server_api import ServerApi
-from . import log_module
-import datetime
-
-MONGO_URL = os.environ.get("MONGO_URL", None)
-MONGO_PWD = os.environ.get("MONGO_PWD", None)
-MONGO_USR = os.environ.get("MONGO_USR", None)
-uri = f"mongodb+srv://{MONGO_USR}:{MONGO_PWD}@{MONGO_URL}/?retryWrites=true&w=majority"
-uri = f"mongodb://localhost:27017"
-
-tz = datetime.timezone(datetime.timedelta(hours=-4))
-
-def check_user(data, create=True):
-    client = MongoClient(uri, server_api=ServerApi('1'))
-    user_table = client.ChatDB.users
-    user = user_table.find_one({"gid":data["gid"]})
-    if not user and create:
-        user = {
-            "name": data["name"],
-            "tokens": {},
-            "approved": False,
-            "created": datetime.datetime.now(tz),
-            "description": "",
-            "email": data["email"],
-            "gid": data["gid"],
-            "role": "on hold"
-        }
-        user_table.insert_one(user)
-    user.pop("_id", "")
-    return user
-
-def add_description(gid, message):
-    client = MongoClient(uri, server_api=ServerApi('1'))
-    user_table = client.ChatDB.users
-    sel = {"gid":gid}
-    upd = {"$set": { "description": message}}
-    user_table.update_one(sel, upd)
-
-        

modules/model.py

@@ -0,0 +1,66 @@
+import os 
+from pymongo.mongo_client import MongoClient
+from pymongo.server_api import ServerApi
+from . import log_module, security
+from datetime import timezone, datetime, timedelta
+from pydantic import BaseModel
+
+
+MONGO_URL = os.environ.get("MONGO_URL", None)
+MONGO_PWD = os.environ.get("MONGO_PWD", None)
+MONGO_USR = os.environ.get("MONGO_USR", None)
+uri = f"mongodb+srv://{MONGO_USR}:{MONGO_PWD}@{MONGO_URL}/?retryWrites=true&w=majority"
+#uri = f"mongodb://localhost:27017"
+client = MongoClient(uri, server_api=ServerApi('1'))
+user_table = client.ChatDB.users
+
+tz = timezone(timedelta(hours=-4))
+
+class User(BaseModel):
+    name: str
+    tokens: dict = {}
+    created: datetime = datetime.now(tz)
+    approved: datetime | None = None
+    description: str = ""
+    email: str
+    gid: str
+    role: str = "on hold"
+
+    @classmethod
+    def find(cls, gid: str):
+        found = user_table.find_one({"gid":gid})
+        if found:
+            return cls(**found)
+        return None
+
+    @classmethod
+    def find_from_token(cls, token: str):
+        data = security.validate_token(token)
+        user = user_table.find_one({"gid":data["gid"]})
+        return cls(**user)
+        
+        
+    @classmethod
+    def find_or_create(cls, data: dict):
+        found = cls.find(data["gid"])
+        if found:
+            return found
+        user = User(**data)
+        user_table.insert_one(user.model_dump())    
+        log_module.logger.info(f"User created: {user.gid}")
+        return user
+    
+    def create_token(self):
+        return security.create_jwt_token({"gid": self.gid})
+    
+    def update_description(self, message):
+        log_module.logger.info(f"Description Updated: {self.gid}")
+        user_table.update_one(
+            {"gid":self.gid}, 
+            {"$set": { "description": message}}
+        )
+        self.description = message
+
+    def can_use(self, activity):
+        return security.can_use(self.role, activity)
+        

modules/security.py

@@ -7,8 +7,6 @@ from . import log_module
 from datetime import datetime, timedelta
 
 
-
-
 users = json.loads(str(os.getenv("USER_KEYS")).replace("\n", ""))
 for key in users:
     if key == "master": continue
@@ -19,45 +17,34 @@ JWT_SECRET = users["master"]
 JWT_ALGORITHM = "HS256"
 JWT_EXPIRATION_TIME_MINUTES = 30
 
-async def authenticate_user(credentials: HTTPBasicCredentials) -> bool:
-    RedirectResponse(url='/login')
-    return True
+credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+
 
 def create_jwt_token(data):
-    to_encode = {"data": data}
+    to_encode = data
     expire = datetime.utcnow() + timedelta(minutes=JWT_EXPIRATION_TIME_MINUTES)
     to_encode.update({"exp": expire})
     encoded_jwt = jwt.encode(to_encode, JWT_SECRET, algorithm=JWT_ALGORITHM)
     return encoded_jwt
 
-async def validate_cookie(request: Request):
+def validate_jwt_token(token):
     try:
-        token = request.cookies.get('token')
         payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
-        data = {
-            "token": create_jwt_token(payload["data"]),
-            "token_data": payload["data"]
-        }
-        return (data, False)
+        return payload
     except Exception as e: 
-        log_module.logger.error(repr(e) + " - " + str(data))
-        return ({}, RedirectResponse(url='/login'))
-    
-async def validate_token(request: Request):
-    data = {}
-    try: 
-        data = await request.json()
-        token = data.pop("token")
-        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
-        data["token_data"] = payload["data"]
-    except Exception as e: 
-        log_module.logger.error(repr(e) + " - " + str(data))
+        log_module.logger.error(repr(e) + " - Invalid token: " + str(token))
         raise HTTPException(status_code=404, detail="Token inválido")    
-    return data
-
-def user_scale(role, action):
-    actions = {
-        "chat": ["user", "admin"],
-        "read_log": ["admin"]
-    }
-    return role in actions[action]
+
+def token_from_cookie(request:Request):
+    if token := request.cookies.get('token', None):
+        return validate_jwt_token(token)
+    return None
+
+def token_from_json(request:Request):
+    if token := request.headers.get("Autorization") :
+        return validate_jwt_token(token)
+    return None   

modules/token_counter.py

@@ -0,0 +1,39 @@
+def num_tokens_from_messages(messages, model="gpt-3.5-turbo-0613"):
+    """Return the number of tokens used by a list of messages."""
+    try:
+        encoding = tiktoken.encoding_for_model(model)
+    except KeyError:
+        print("Warning: model not found. Using cl100k_base encoding.")
+        encoding = tiktoken.get_encoding("cl100k_base")
+    if model in {
+        "gpt-3.5-turbo-0613",
+        "gpt-3.5-turbo-16k-0613",
+        "gpt-4-0314",
+        "gpt-4-32k-0314",
+        "gpt-4-0613",
+        "gpt-4-32k-0613",
+        }:
+        tokens_per_message = 3
+        tokens_per_name = 1
+    elif model == "gpt-3.5-turbo-0301":
+        tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
+        tokens_per_name = -1  # if there's a name, the role is omitted
+    elif "gpt-3.5-turbo" in model:
+        print("Warning: gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
+        return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0613")
+    elif "gpt-4" in model:
+        print("Warning: gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
+        return num_tokens_from_messages(messages, model="gpt-4-0613")
+    else:
+        raise NotImplementedError(
+            f"""num_tokens_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens."""
+        )
+    num_tokens = 0
+    for message in messages:
+        num_tokens += tokens_per_message
+        for key, value in message.items():
+            num_tokens += len(encoding.encode(value))
+            if key == "name":
+                num_tokens += tokens_per_name
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
+    return num_tokens