-
Notifications
You must be signed in to change notification settings - Fork 7
/
main.py
127 lines (86 loc) · 2.57 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import logging
import os
from typing import List, Optional
import sentry_sdk
from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from ai import ai_chat_thread
from pipeline import MaxPipeline, Entries
from slack import app as slack_app
load_dotenv() # take environment variables from .env.
sentry_sdk.init(
dsn="https://4a3780ef52824c52b13eeab44ea73a14@o1015702.ingest.sentry.io/4505009495605248",
# Set traces_sample_rate to 1.0 to capture 100%
# of transactions for performance monitoring.
# We recommend adjusting this value in production,
traces_sample_rate=1.0,
)
origins = [
"http://localhost",
"http://localhost:8000",
"http://localhost:8001",
"http://localhost:8002",
"https://app.posthog.com",
"https://posthog.com",
]
# FastAPI App
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
class Message(BaseModel):
role: str
content: str
class Query(BaseModel):
query: str
class GitHubRepo(BaseModel):
repo: Optional[str]
pipeline = MaxPipeline(openai_token=os.getenv("OPENAI_TOKEN"))
@app.post("/entries")
def create_entries(entries: Entries):
pipeline.embed_markdown_document(entries)
return []
@app.post("/_git")
def create_git_entries(gh_repo: GitHubRepo):
pipeline.embed_git_repo(gh_repo=gh_repo.repo)
return {"status": "ok"}
@app.post("/_chat")
def test_chat(query: Query):
return pipeline.chat(query.query)
@app.post("/_context")
def test_context(query: Query):
return pipeline.retrieve_context(query.query)
@app.post("/spawn")
def receive_spawn():
print("Spawned")
return []
@app.post("/update")
def update_oncall():
return "nope"
@app.post("/chat")
async def chat(messages: List[Message]):
msgs = [msg.dict() for msg in messages]
response = await ai_chat_thread(msgs)
return response
@app.get("/_health")
def health():
return {"status": "ok"}
# Slack Bolt App
from slack_bolt.adapter.fastapi.async_handler import AsyncSlackRequestHandler
app_handler = AsyncSlackRequestHandler(slack_app)
@app.post("/slack/events")
async def slack_events(req: Request):
return await app_handler.handle(req)
@app.get("/slack/oauth_redirect")
async def oauth_redirect(req: Request):
logging.info("Installation completed.")
return await app_handler.handle(req)
@app.get("/slack/install")
async def install(req: Request):
return await app_handler.handle(req)