file update

.gitignore

@@ -1 +1,5 @@
-.venv
+.venv
+cache
+users.db
+.env
+__pycache__

auth.py

@@ -3,61 +3,93 @@ from fastapi.security import OAuth2PasswordBearer
 from jose import JWTError, jwt
 from passlib.context import CryptContext
 from datetime import datetime, timedelta
+from sqlalchemy.orm import Session
+from database import engine, get_db
+from models import Base, User
 import os
 
-# Load secrets from environment variables
-SECRET_KEY = os.getenv("SECRET_KEY")
-ALGORITHM = "HS512"
+Base.metadata.create_all(bind=engine)
+# Load secrets from environment variables or set defaults
+SECRET_KEY = os.getenv("SECRET_KEY", "def6nQHONW99pOPyba9DShny6FB1CJJBigZault")
+ALGORITHM = "HS256"
 ACCESS_TOKEN_EXPIRE_MINUTES = 30
+REFRESH_TOKEN_EXPIRE_DAYS = 7
 
-# Password hashing
+# OAuth2 scheme (Ensure the token URL matches the actual login endpoint)
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
+
+# Password hashing context
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
-# OAuth2 scheme
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+def hash_password(password: str) -> str:
+    return pwd_context.hash(password)
+
+def register_user(username: str, password: str, db: Session):
+    existing_user = db.query(User).filter(User.username == username).first()
+    if existing_user:
+        raise HTTPException(status_code=400, detail="Username already taken")
 
-# Mock user database (replace with a real database in production)
-fake_users_db = {
-    "admin": {
-        "username": "admin",
-        "hashed_password": pwd_context.hash(os.getenv("ADMIN_PASSWORD")),
-    }
-}
+    hashed_password = hash_password(password)
+    new_user = User(username=username, password=hashed_password)
+    
+    db.add(new_user)
+    db.commit()
+    db.refresh(new_user)
+    
+    return new_user
 
-def verify_password(plain_password, hashed_password):
+def verify_password(plain_password: str, hashed_password: str) -> bool:
     return pwd_context.verify(plain_password, hashed_password)
 
-def authenticate_user(username: str, password: str):
-    user = fake_users_db.get(username)
-    if not user or not verify_password(password, user["hashed_password"]):
-        return False
+def authenticate_user(username: str, password: str, db: Session):
+    user = db.query(User).filter(User.username == username).first()
+    if not user or not verify_password(password, user.password):
+        return None
     return user
 
-def create_access_token(data: dict, expires_delta: timedelta = None):
+def create_token(data: dict, expires_delta: timedelta, secret_key: str) -> str:
+    """Generate JWT token."""
     to_encode = data.copy()
-    if expires_delta:
-        expire = datetime.utcnow() + expires_delta
-    else:
-        expire = datetime.utcnow() + timedelta(minutes=15)
+    expire = datetime.utcnow() + expires_delta
     to_encode.update({"exp": expire})
-    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    
+    encoded_jwt = jwt.encode(to_encode, secret_key, algorithm=ALGORITHM)
     return encoded_jwt
 
-async def get_current_user(token: str = Depends(oauth2_scheme)):
-    credentials_exception = HTTPException(
-        status_code=status.HTTP_401_UNAUTHORIZED,
-        detail="Could not validate credentials",
-        headers={"WWW-Authenticate": "Bearer"},
-    )
+def verify_token(token: str, secret_key: str) -> str:
+    """Verifies JWT token and extracts the username."""
     try:
-        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        payload = jwt.decode(token, secret_key, algorithms=[ALGORITHM])
         username: str = payload.get("sub")
         if username is None:
-            raise credentials_exception
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid token",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        return username
     except JWTError:
-        raise credentials_exception
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        ) 
 
-    user = fake_users_db.get(username)
-    if user is None:
-        raise credentials_exception
-    return user
+def verify_access_token(token: str = Depends(oauth2_scheme)) -> str:
+    """Verifies access token and returns the username."""
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username: str = payload.get("sub")
+        if username is None:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid access token",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        return username
+    except JWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )

cache.py

@@ -0,0 +1,23 @@
+import os
+from typing import List, Optional
+from sentence_transformers import SentenceTransformer
+from functools import lru_cache
+from dotenv import load_dotenv
+
+@lru_cache(maxsize=1)
+def get_sentence_transformer() -> SentenceTransformer:
+    """Loads and caches the Sentence Transformer model."""
+    try:
+        model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+        print("✅ Sentence Transformer Model Loaded")
+        return model
+    except Exception as e:
+        print(f"❌ Error loading Sentence Transformer: {str(e)}")
+        raise RuntimeError("Failed to load Sentence Transformer model.")
+
+    
+def get_cached_embeddings(text: str, model_type: str) -> Optional[List[float]]:
+    pass
+
+def set_cached_embeddings(text: str, model_type: str, embeddings: List[float]):
+    pass

database.py

@@ -0,0 +1,16 @@
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+DATABASE_URL = "sqlite:///./users.db" 
+
+engine = create_engine(DATABASE_URL)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+Base = declarative_base()
+
+def get_db():
+    """Dependency function to get the database session."""
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()

gen.py

@@ -1,13 +0,0 @@
-from fastapi import WebSocket
-from transformers import pipeline
-import asyncio
-
-# Load the model
-model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
-generator = pipeline("text-generation", model=model_name)
-
-async def generate_text_stream(prompt: str, websocket: WebSocket):
-    for i in range(10):  # Simulate streaming (replace with actual model inference)
-        chunk = generator(prompt, max_length=i + 10, do_sample=True)[0]["generated_text"]
-        await websocket.send_text(chunk)
-        await asyncio.sleep(0.1)  # Simulate delay

main.py

@@ -1,44 +1,118 @@
-from fastapi import FastAPI, WebSocket, Depends, HTTPException
-from auth import get_current_user, authenticate_user, create_access_token
-from gen import generate_text_stream
+from fastapi import FastAPI, Depends, HTTPException, status
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from jose import JWTError
+from schemas import UserRegister, TokenResponse, RefreshTokenRequest, QueryInput
+from auth import  register_user, get_db, authenticate_user, create_token, verify_token, verify_access_token, Session
+from utils import search_with_llm
 from fastapi.security import OAuth2PasswordRequestForm
-from fastapi.middleware.cors import CORSMiddleware
-from middleware import setup_rate_limiter
+from pathlib import Path
+from datetime import timedelta
 import os
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+SECRET_KEY = os.getenv("SECRET_KEY", 'def6nQHONW99pOPyba9DShny6FB1CJJBigZault')
+REFRESH_SECRET_KEY = os.getenv("REFRESH_SECRET_KEY", SECRET_KEY)
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 30
+REFRESH_TOKEN_EXPIRE_DAYS = 7
 
 app = FastAPI()
 
+# Entry Endpoint 
+@app.get('/')
+def index() -> FileResponse:
+    file_path = Path(__file__).parent / 'static' / 'index.html'
+    return FileResponse(path=str(file_path), media_type='text/html')
 
-# Apply rate limiting middleware
-setup_rate_limiter(app)
+@app.post("/register")
+async def register(user: UserRegister, db: Session = Depends(get_db)):
+    """Registers a new user."""
+    new_user = register_user(user.username, user.password, db)
+    return {"message": "User registered successfully", "user": new_user.username}
 
-# CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+@app.post("/login", response_model=TokenResponse)
+async def login(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
+    try:
+        user = authenticate_user(form_data.username, form_data.password, db)
 
-# Login endpoint
-@app.post("/token")
-async def login(form_data: OAuth2PasswordRequestForm = Depends()):
-    user = authenticate_user(form_data.username, form_data.password)
-    if not user:
-        raise HTTPException(status_code=400, detail="Incorrect username or password")
-    access_token = create_access_token(data={"sub": user["username"]})
-    return {"access_token": access_token, "token_type": "bearer"}
+        if not user:
+           raise HTTPException(
+               status_code=status.HTTP_401_UNAUTHORIZED,
+               detail="Invalid username or password",
+               headers={"WWW-Authenticate": "Bearer"},
+           )
 
-# WebSocket endpoint for streaming
-@app.websocket("/generate")
-async def websocket_generate(websocket: WebSocket, token: str):
-    await websocket.accept()
+        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+        refresh_token_expires = timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS)
+        
+        access_token = create_token(
+            data={"sub": user.username}, 
+            expires_delta=access_token_expires, 
+            secret_key=SECRET_KEY
+        )
+        refresh_token = create_token(
+            data={"sub": user.username}, 
+            expires_delta=refresh_token_expires, 
+            secret_key=REFRESH_SECRET_KEY
+        )
+
+        return {
+            "access_token": access_token, 
+            "refresh_token": refresh_token, 
+            "token_type": "bearer"
+        }
+    except Exception as e:
+        logging.error(f"Login error: {str(e)}")
+        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Internal server error")
+
+@app.post("/refresh", response_model=TokenResponse)
+async def refresh(refresh_request: RefreshTokenRequest):
     try:
-        user = get_current_user(token)
-        prompt = await websocket.receive_text()
-        await generate_text_stream(prompt, websocket)
+        # Verify the refresh token
+        username = verify_token(refresh_request.refresh_token, REFRESH_SECRET_KEY)
+        
+        # Create new access token
+        access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+        access_token = create_token(
+            data={"sub": username},
+            expires_delta=access_token_expires,
+            secret_key=SECRET_KEY
+        )
+        
+        return {
+            "access_token": access_token,
+            "refresh_token": refresh_request.refresh_token,  # Return the same refresh token
+            "token_type": "bearer"
+        }
+        
+    except JWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    
+
+
+@app.post("/search")
+async def search(
+    query_input: QueryInput,
+    username: str = Depends(verify_access_token),
+):
+    try:
+        response = search_with_llm(query_input.query)
+        return {"response": response}
     except Exception as e:
-        await websocket.send_text(f"Error: {str(e)}")
-    finally:
-        await websocket.close()
+        raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}")
+
+# WebSocket endpoint for streaming
+@app.on_event("startup")
+async def startup_event():
+    os.makedirs("./cache", exist_ok=True)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

middleware.py

@@ -1,18 +0,0 @@
-from fastapi import Request
-from slowapi import Limiter
-from slowapi.util import get_remote_address
-from slowapi.middleware import SlowAPIMiddleware
-
-
-limiter = Limiter(key_func=get_remote_address)
-
-async def log_requests(request: Request, call_next):
-    print(f"Request: {request.method} {request.url}")
-    response = await call_next(request)
-    print(f"Response: {response.status_code}")
-    return response
-
-# Apply the rate limiter middleware
-def setup_rate_limiter(app):
-    app.state.limiter = limiter
-    app.add_middleware(SlowAPIMiddleware)

models.py

@@ -0,0 +1,10 @@
+from sqlalchemy import Column, Integer, String
+from database import Base
+
+class User(Base):
+    """SQLAlchemy model for storing users."""
+    __tablename__ = "users"
+
+    id = Column(Integer, primary_key=True, index=True)
+    username = Column(String, unique=True, index=True)
+    password = Column(String)  # Hashed password

requirements.txt

@@ -22,6 +22,17 @@ loguru  # Optional: for better logging
 # Testing (optional for development)
 pytest
 httpx  # For async HTTP requests in tests
+bcrypt
+openai
+datasets
+diskcache
+sentence_transformers
+pinecone-client
+python-multipart
+sqlalchemy
+psycopg2
+python-dotenv
+ollama
 
 # Docker and deployment
 gunicorn  # Optional: for production deployment

schemas.py

@@ -1,49 +1,16 @@
-from pydantic import BaseModel, EmailStr
+from pydantic import BaseModel, Field
 
-# Authentication
-class Token(BaseModel):
-    access_token: str
-    token_type: str
-
-class TokenData(BaseModel):
-    username: str | None = None
-
-class User(BaseModel):
-    username: str
-    email: EmailStr | None = None
-
-class UserInDB(User):
-    hashed_password: str
-
-class LoginRequest(BaseModel):
+class UserRegister(BaseModel):
     username: str
     password: str
 
-# Generation
-class GenerationRequest(BaseModel):
-    prompt: str
-    max_length: int = 100
-    temperature: float = 0.7
-    top_k: int = 50
-    top_p: float = 0.95
-
-class GenerationResponse(BaseModel):
-    generated_text: str
-
-# WebSocket
-class WebSocketMessage(BaseModel):
-    prompt: str
+class QueryInput(BaseModel):
+    query: str 
 
-# Error Handling
-class HTTPError(BaseModel):
-    detail: str
-
-    class Config:
-        schema_extra = {
-            "example": {"detail": "Error message"},
-        }
+class TokenResponse(BaseModel):
+    access_token: str
+    token_type: str
+    refresh_token: str
 
-# Rate Limiting
-class RateLimitResponse(BaseModel):
-    message: str
-    retry_after: int
+class RefreshTokenRequest(BaseModel):
+    refresh_token: str

utils.py

@@ -0,0 +1,21 @@
+import os
+import ollama
+from typing import List, Dict
+
+def cosine_similarity(embedding_0, embedding_1):
+    pass
+
+
+def generate_embedding(model, text: str, model_type: str) -> List[float]:
+    pass
+
+def search_with_llm(query: str, model: str = "llama3.2"):
+    try:
+        response = ollama.chat(
+            model=model,
+            messages=[{"role": "user", "content": query}]
+        )
+        return response["message"]["content"]
+    except Exception as e:
+        return f"❌ Error processing request: {str(e)}"
+