Stay Hungry,Stay Foolish!

ai-assistant stream with websocket.

ai-assistant

stream with websocket.

 

https://github.com/dpills/ai-assistant/blob/master/main.py

from typing import AsyncGenerator, NoReturn

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from openai import AsyncOpenAI

load_dotenv()

app = FastAPI()
client = AsyncOpenAI()

with open("index.html") as f:
    html = f.read()


async def get_ai_response(message: str) -> AsyncGenerator[str, None]:
    """
    OpenAI Response
    """
    response = await client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant, skilled in explaining "
                    "complex concepts in simple terms."
                ),
            },
            {
                "role": "user",
                "content": message,
            },
        ],
        stream=True,
    )

    all_content = ""
    async for chunk in response:
        content = chunk.choices[0].delta.content
        if content:
            all_content += content
            yield all_content


@app.get("/")
async def web_app() -> HTMLResponse:
    """
    Web App
    """
    return HTMLResponse(html)


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> NoReturn:
    """
    Websocket for AI responses
    """
    await websocket.accept()
    while True:
        message = await websocket.receive_text()
        async for text in get_ai_response(message):
            await websocket.send_text(text)


if __name__ == "__main__":
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        log_level="debug",
        reload=True,
    )

 

 

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="utf-8" />
    <title>AI Assistant 🤓</title>
    <meta name="viewport" content="initial-scale=1, width=device-width" />
    <script src="https://unpkg.com/react@latest/umd/react.development.js" crossorigin="anonymous"></script>
    <script src="https://unpkg.com/react-dom@latest/umd/react-dom.development.js"></script>
    <script src="https://unpkg.com/@mui/material@latest/umd/material-ui.development.js"
        crossorigin="anonymous"></script>
    <script src="https://unpkg.com/@babel/standalone@latest/babel.min.js" crossorigin="anonymous"></script>
    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
    <link rel="preconnect" href="https://fonts.googleapis.com" />
    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
    <link rel="stylesheet"
        href="https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;500;600;700&display=swap" />
    <link rel="stylesheet" href="https://fonts.googleapis.com/icon?family=Material+Icons" />
</head>

<body>
    <div id="root"></div>
    <script type="text/babel">
        const {
            colors,
            CssBaseline,
            ThemeProvider,
            Typography,
            TextField,
            Container,
            createTheme,
            Box,
            Skeleton,
        } = MaterialUI;

        const theme = createTheme({
            palette: {
                mode: 'dark'
            },
        });
        const WS = new WebSocket("ws://localhost:8000/ws");

        function App() {
            const [response, setResponse] = React.useState("");
            const [question, setQuestion] = React.useState("");
            const [loading, setLoading] = React.useState(false);

            React.useEffect(() => {
                WS.onmessage = (event) => {
                    setLoading(false);
                    setResponse(marked.parse(event.data));
                };
            }, []);

            return (
                <Container maxWidth="lg">
                    <Box sx={{ my: 4 }}>
                        <Typography variant="h4" component="h1" gutterBottom>
                            AI Assistant 🤓
                        </Typography>
                        <TextField
                            id="outlined-basic"
                            label="Ask me Anything"
                            variant="outlined"
                            style={{ width: '100%' }}
                            value={question}
                            disabled={loading}
                            onChange={e => {
                                setQuestion(e.target.value)
                            }}
                            onKeyUp={e => {
                                setLoading(false)
                                if (e.key === "Enter") {
                                    setResponse('')
                                    setLoading(true)
                                    WS.send(question);
                                }
                            }}
                        />
                    </Box>
                    {!response && loading && (<>
                        <Skeleton />
                        <Skeleton animation="wave" />
                        <Skeleton animation={false} /></>)}
                    {response && <Typography dangerouslySetInnerHTML={{ __html: response }} />}
                </Container>
            );
        }

        ReactDOM.createRoot(document.getElementById('root')).render(
            <ThemeProvider theme={theme}>
                <CssBaseline />
                <App />
            </ThemeProvider>,
        );
    </script>
</body>

</html>

 

posted @ 2024-12-04 22:57  lightsong  阅读(3)  评论(0)  编辑  收藏  举报
Life Is Short, We Need Ship To Travel