2 changed files with 77 additions and 0 deletions
@ -0,0 +1,38 @@ |
|||
import time |
|||
from typing import Annotated |
|||
|
|||
from fastapi import Depends, FastAPI, HTTPException |
|||
from fastapi.responses import StreamingResponse |
|||
from sqlmodel import Field, Session, SQLModel, create_engine |
|||
|
|||
engine = create_engine("postgresql+psycopg://postgres:postgres@localhost/db") |
|||
|
|||
|
|||
class User(SQLModel, table=True): |
|||
id: int | None = Field(default=None, primary_key=True) |
|||
name: str |
|||
|
|||
|
|||
app = FastAPI() |
|||
|
|||
|
|||
def get_session(): |
|||
with Session(engine) as session: |
|||
yield session |
|||
|
|||
|
|||
def get_user(user_id: int, session: Annotated[Session, Depends(get_session)]): |
|||
user = session.get(User, user_id) |
|||
if not user: |
|||
raise HTTPException(status_code=403, detail="Not authorized") |
|||
|
|||
|
|||
def generate_stream(query: str): |
|||
for ch in query: |
|||
yield ch |
|||
time.sleep(0.1) |
|||
|
|||
|
|||
@app.get("/generate", dependencies=[Depends(get_user)]) |
|||
def generate(query: str): |
|||
return StreamingResponse(content=generate_stream(query)) |
|||
@ -0,0 +1,39 @@ |
|||
import time |
|||
from typing import Annotated |
|||
|
|||
from fastapi import Depends, FastAPI, HTTPException |
|||
from fastapi.responses import StreamingResponse |
|||
from sqlmodel import Field, Session, SQLModel, create_engine |
|||
|
|||
engine = create_engine("postgresql+psycopg://postgres:postgres@localhost/db") |
|||
|
|||
|
|||
class User(SQLModel, table=True): |
|||
id: int | None = Field(default=None, primary_key=True) |
|||
name: str |
|||
|
|||
|
|||
app = FastAPI() |
|||
|
|||
|
|||
def get_session(): |
|||
with Session(engine) as session: |
|||
yield session |
|||
|
|||
|
|||
def get_user(user_id: int, session: Annotated[Session, Depends(get_session)]): |
|||
user = session.get(User, user_id) |
|||
if not user: |
|||
raise HTTPException(status_code=403, detail="Not authorized") |
|||
session.close() |
|||
|
|||
|
|||
def generate_stream(query: str): |
|||
for ch in query: |
|||
yield ch |
|||
time.sleep(0.1) |
|||
|
|||
|
|||
@app.get("/generate", dependencies=[Depends(get_user)]) |
|||
def generate(query: str): |
|||
return StreamingResponse(content=generate_stream(query)) |
|||
Loading…
Reference in new issue