Spaces:
Paused
Paused
| import discord | |
| from discord.ext import commands | |
| from discord.ext.commands import Context | |
| from fastapi import FastAPI, HTTPException, Request | |
| from fastapi.responses import JSONResponse | |
| import gradio as gr | |
| # Initialize FastAPI app | |
| app = FastAPI() | |
| # Initialize Gradio Llama model | |
| llm = gr.Llama(model_path="model.gguf", n_ctx=4000, n_threads=2, chat_format="chatml") | |
| # Initialize Discord bot | |
| bot = commands.Bot(command_prefix='&') # Define the command prefix | |
| # Global variable to store the channel where chats will be sent | |
| chat_channel = None | |
| # Define the command to set the chat channel | |
| async def set_channel(ctx: Context, channel: discord.TextChannel): | |
| global chat_channel | |
| chat_channel = channel | |
| await ctx.send(f"Chat channel set to {channel.mention}") | |
| # Define the function to handle the chat endpoint | |
| async def chat_post(request: Request): | |
| global chat_channel | |
| if chat_channel is None: | |
| raise HTTPException(status_code=400, detail="Chat channel is not set") | |
| data = await request.json() | |
| message = data.get("message") | |
| history = data.get("history", []) | |
| temperature = data.get("temperature", 0.3) | |
| max_tokens = data.get("max_tokens", 512) | |
| async def generate(): | |
| system_prompt = "You are OpenChat, a useful AI assistant." | |
| formatted_prompt = [{"role": "system", "content": system_prompt}] | |
| for user_prompt, bot_response in history: | |
| formatted_prompt.append({"role": "user", "content": user_prompt}) | |
| formatted_prompt.append({"role": "assistant", "content": bot_response}) | |
| formatted_prompt.append({"role": "user", "content": message}) | |
| stream_response = llm.create_chat_completion(messages=formatted_prompt, temperature=temperature, | |
| max_tokens=max_tokens, stream=True) | |
| response = "" | |
| for chunk in stream_response: | |
| if len(chunk['choices'][0]["delta"]) != 0 and "content" in chunk['choices'][0]["delta"]: | |
| response += chunk['choices'][0]["delta"]["content"] | |
| yield response | |
| # Send the generated response to the chat channel | |
| async for response in generate(): | |
| await chat_channel.send(response) | |
| return JSONResponse(content={"response": "Message sent to chat channel"}) | |
| # Define the function to handle the GET request for chat | |
| async def chat_get(): | |
| return {"message": "Send a POST request to this endpoint to chat."} | |
| # Run the bot | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=8000) | |
| bot.run('YOUR_DISCORD_BOT_TOKEN') # Replace 'YOUR_DISCORD_BOT_TOKEN' with your actual bot token | |