nick1221/system1

1from fastapi import FastAPI, Request2from fastapi.middleware.cors import CORSMiddleware3from fastapi.responses import JSONResponse, StreamingResponse4from fastapi.encoders import jsonable_encoder5from typing import Optional, List, Literal6import requests7from pydantic import BaseModel, Field8import json9import os10from system1 import System111from dotenv import load_dotenv12import time13load_dotenv()14import uvicorn15import asyncio16import shortuuid17 18 19vector = System1(llm_provider="openai", model_id="gpt-4o", agent_type="openai_functions")20app = FastAPI()21 22origins = [23    "http://localhost:5173",24    "localhost:8000",25    "localhost:8080"26    "*"27]28 29# Add CORS middleware30app.add_middleware(31    CORSMiddleware,32    allow_origins=origins,33    allow_credentials=True,34    allow_methods=["*"],35    allow_headers=["*"],36)37 38 39@app.get('/')40def index():41    return {"message": "Vector API"}42 43 44@app.get('/search')45async def vector_search(query: str):46    message = vector.query(query)47    return JSONResponse(content=jsonable_encoder({"message": message}))48 49@app.get('/chat')50async def vector_chat(query: str):51    message = vector.chat(query)52    return JSONResponse(content=jsonable_encoder({"message": message}))53 54 55from fastapi import Request56 57class Message(BaseModel):58    role: str59    content: str60 61class CompletionRequest(BaseModel):62    model: str63    messages: List[Message]64    max_tokens: int = None65    temperature: float = 1.066    stream: bool = False67 68class DeltaMessage(BaseModel):69    role: Optional[str] = None70    content: Optional[str] = None71 72 73class ChatCompletionResponseStreamChoice(BaseModel):74    index: int75    delta: DeltaMessage76    finish_reason: Optional[Literal["stop", "length"]] = None77 78 79class ChatCompletionStreamResponse(BaseModel):80    id: str = Field(default_factory=lambda: f"chatcmpl-{shortuuid.random()}")81    object: str = "chat.completion.chunk"82    created: int = Field(default_factory=lambda: int(time.time()))83    model: str84 85@app.post("/v1/chat/completions")86async def completions(request: CompletionRequest):87    if request.stream:88        print("STREAM")89        async def stream_response():90            for i in range(1, 11):91                delta = {92                    "id": "msg_123",93                    "object": "thread.message.delta",94                    "delta": {95                        "content": [96                            {97                                "index": 0,98                                "type": "text",99                                "text": { "value": f"Stream {i}...", "annotations": [] }100                            }101                        ]102                    }103                }104                yield json.dumps(delta)105            # Final delta to indicate the end of the stream106            final_delta = {107                "id": "msg_123",108                "object": "thread.message.delta",109                "delta": {110                    "content": [111                        {112                            "index": 0,113                            "type": "text",114                            "text": { "value": "", "annotations": [] }115                        }116                    ]117                }118            }119            yield json.dumps(final_delta)120        121        return StreamingResponse(stream_response(), media_type="application/json")122    else:123        # Existing functionality for non-streaming responses124        response = vector.chat_messages(request.messages)125        prompt_tokens = sum(len(message.content) for message in request.messages)126        completion_tokens = len(str(response))127        total_tokens = prompt_tokens + completion_tokens128 129        formatted_response = {130            "id": "vector-123",131            "object": "chat.completion",132            "created": int(time.time()),133            "model": "vector-openai",134            "system_fingerprint": "fp_44709d6fcb",135            "choices": [{136                "index": 0,137                "message": {138                    "role": "assistant",139                    "content": response,140                },141                "logprobs": None,142                "finish_reason": "stop"143            }],144            "usage": {145                "prompt_tokens": prompt_tokens,146                "completion_tokens": completion_tokens,147                "total_tokens": total_tokens148            }149        }150        return jsonable_encoder(formatted_response)151 152 153if __name__ == '__main__':154    uvicorn.run("app:app", host='127.0.0.1', port=int(os.environ.get('PORT', '8080')), reload=True)