Skip to content

Commit 8998b44

Browse files
use black formatter
1 parent 57b5adb commit 8998b44

File tree

6 files changed

+132
-97
lines changed

6 files changed

+132
-97
lines changed

backend/app/main.py

+4-6
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,7 @@
1313
app = FastAPI()
1414

1515

16-
origins = [
17-
"http://localhost:3000",
18-
"https://gitdiagram.com"
19-
]
16+
origins = ["http://localhost:3000", "https://gitdiagram.com"]
2017

2118
app.add_middleware(
2219
CORSMiddleware,
@@ -31,8 +28,9 @@
3128
app.add_middleware(Analytics, api_key=API_ANALYTICS_KEY)
3229

3330
app.state.limiter = limiter
34-
app.add_exception_handler(RateLimitExceeded, cast(
35-
ExceptionMiddleware, _rate_limit_exceeded_handler))
31+
app.add_exception_handler(
32+
RateLimitExceeded, cast(ExceptionMiddleware, _rate_limit_exceeded_handler)
33+
)
3634

3735
app.include_router(generate.router)
3836
app.include_router(modify.router)

backend/app/routers/generate.py

+33-32
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@
33
from app.services.github_service import GitHubService
44
from app.services.claude_service import ClaudeService
55
from app.core.limiter import limiter
6-
import os
7-
from app.prompts import SYSTEM_FIRST_PROMPT, SYSTEM_SECOND_PROMPT, SYSTEM_THIRD_PROMPT, ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
6+
from app.prompts import (
7+
SYSTEM_FIRST_PROMPT,
8+
SYSTEM_SECOND_PROMPT,
9+
SYSTEM_THIRD_PROMPT,
10+
ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT,
11+
)
812
from anthropic._exceptions import RateLimitError
913
from pydantic import BaseModel
1014
from functools import lru_cache
@@ -29,11 +33,7 @@ def get_cached_github_data(username: str, repo: str):
2933
file_tree = github_service.get_github_file_paths_as_list(username, repo)
3034
readme = github_service.get_github_readme(username, repo)
3135

32-
return {
33-
"default_branch": default_branch,
34-
"file_tree": file_tree,
35-
"readme": readme
36-
}
36+
return {"default_branch": default_branch, "file_tree": file_tree, "readme": readme}
3737

3838

3939
class ApiRequest(BaseModel):
@@ -51,7 +51,13 @@ async def generate(request: Request, body: ApiRequest):
5151
if len(body.instructions) > 1000:
5252
return {"error": "Instructions exceed maximum length of 1000 characters"}
5353

54-
if body.repo in ["fastapi", "streamlit", "flask", "api-analytics", "monkeytype"]:
54+
if body.repo in [
55+
"fastapi",
56+
"streamlit",
57+
"flask",
58+
"api-analytics",
59+
"monkeytype",
60+
]:
5561
return {"error": "Example repos cannot be regenerated"}
5662

5763
# Get cached github data
@@ -71,7 +77,7 @@ async def generate(request: Request, body: ApiRequest):
7177
return {
7278
"error": f"File tree and README combined exceeds token limit (50,000). Current size: {token_count} tokens. This GitHub repository is too large for my wallet, but you can continue by providing your own Anthropic API key.",
7379
"token_count": token_count,
74-
"requires_api_key": True
80+
"requires_api_key": True,
7581
}
7682
elif token_count > 200000:
7783
return {
@@ -82,20 +88,22 @@ async def generate(request: Request, body: ApiRequest):
8288
first_system_prompt = SYSTEM_FIRST_PROMPT
8389
third_system_prompt = SYSTEM_THIRD_PROMPT
8490
if body.instructions:
85-
first_system_prompt = first_system_prompt + \
86-
"\n" + ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
87-
third_system_prompt = third_system_prompt + \
88-
"\n" + ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
91+
first_system_prompt = (
92+
first_system_prompt + "\n" + ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
93+
)
94+
third_system_prompt = (
95+
third_system_prompt + "\n" + ADDITIONAL_SYSTEM_INSTRUCTIONS_PROMPT
96+
)
8997

9098
# get the explanation for sysdesign from claude
9199
explanation = claude_service.call_claude_api(
92100
system_prompt=first_system_prompt,
93101
data={
94102
"file_tree": file_tree,
95103
"readme": readme,
96-
"instructions": body.instructions
104+
"instructions": body.instructions,
97105
},
98-
api_key=body.api_key
106+
api_key=body.api_key,
99107
)
100108

101109
# Check for BAD_INSTRUCTIONS response
@@ -104,18 +112,14 @@ async def generate(request: Request, body: ApiRequest):
104112

105113
full_second_response = claude_service.call_claude_api(
106114
system_prompt=SYSTEM_SECOND_PROMPT,
107-
data={
108-
"explanation": explanation,
109-
"file_tree": file_tree
110-
}
115+
data={"explanation": explanation, "file_tree": file_tree},
111116
)
112117

113118
# Extract component mapping from the response
114119
start_tag = "<component_mapping>"
115120
end_tag = "</component_mapping>"
116121
component_mapping_text = full_second_response[
117-
full_second_response.find(start_tag):
118-
full_second_response.find(end_tag)
122+
full_second_response.find(start_tag) : full_second_response.find(end_tag)
119123
]
120124

121125
# get mermaid.js code from claude
@@ -124,8 +128,8 @@ async def generate(request: Request, body: ApiRequest):
124128
data={
125129
"explanation": explanation,
126130
"component_mapping": component_mapping_text,
127-
"instructions": body.instructions
128-
}
131+
"instructions": body.instructions,
132+
},
129133
)
130134

131135
# Check for BAD_INSTRUCTIONS response
@@ -134,18 +138,14 @@ async def generate(request: Request, body: ApiRequest):
134138

135139
# Process click events to include full GitHub URLs
136140
processed_diagram = process_click_events(
137-
mermaid_code,
138-
body.username,
139-
body.repo,
140-
default_branch
141+
mermaid_code, body.username, body.repo, default_branch
141142
)
142143

143-
return {"diagram": processed_diagram,
144-
"explanation": explanation}
144+
return {"diagram": processed_diagram, "explanation": explanation}
145145
except RateLimitError as e:
146146
raise HTTPException(
147147
status_code=429,
148-
detail="Service is currently experiencing high demand. Please try again in a few minutes."
148+
detail="Service is currently experiencing high demand. Please try again in a few minutes.",
149149
)
150150
except Exception as e:
151151
return {"error": str(e)}
@@ -184,12 +184,13 @@ def process_click_events(diagram: str, username: str, repo: str, branch: str) ->
184184
Process click events in Mermaid diagram to include full GitHub URLs.
185185
Detects if path is file or directory and uses appropriate URL format.
186186
"""
187+
187188
def replace_path(match):
188189
# Extract the path from the click event
189-
path = match.group(2).strip('"\'')
190+
path = match.group(2).strip("\"'")
190191

191192
# Determine if path is likely a file (has extension) or directory
192-
is_file = '.' in path.split('/')[-1]
193+
is_file = "." in path.split("/")[-1]
193194

194195
# Construct GitHub URL
195196
base_url = f"https://github.com/{username}/{repo}"

backend/app/routers/modify.py

+13-5
Original file line numberDiff line numberDiff line change
@@ -31,19 +31,27 @@ async def modify(request: Request, body: ModifyRequest):
3131
# Check instructions length
3232
if not body.instructions or not body.current_diagram:
3333
return {"error": "Instructions and/or current diagram are required"}
34-
elif len(body.instructions) > 1000 or len(body.current_diagram) > 100000: # just being safe
34+
elif (
35+
len(body.instructions) > 1000 or len(body.current_diagram) > 100000
36+
): # just being safe
3537
return {"error": "Instructions exceed maximum length of 1000 characters"}
3638

37-
if body.repo in ["fastapi", "streamlit", "flask", "api-analytics", "monkeytype"]:
39+
if body.repo in [
40+
"fastapi",
41+
"streamlit",
42+
"flask",
43+
"api-analytics",
44+
"monkeytype",
45+
]:
3846
return {"error": "Example repos cannot be modified"}
3947

4048
modified_mermaid_code = claude_service.call_claude_api(
4149
system_prompt=SYSTEM_MODIFY_PROMPT,
4250
data={
4351
"instructions": body.instructions,
4452
"explanation": body.explanation,
45-
"diagram": body.current_diagram
46-
}
53+
"diagram": body.current_diagram,
54+
},
4755
)
4856

4957
# Check for BAD_INSTRUCTIONS response
@@ -54,7 +62,7 @@ async def modify(request: Request, body: ModifyRequest):
5462
except RateLimitError as e:
5563
raise HTTPException(
5664
status_code=429,
57-
detail="Service is currently experiencing high demand. Please try again in a few minutes."
65+
detail="Service is currently experiencing high demand. Please try again in a few minutes.",
5866
)
5967
except Exception as e:
6068
return {"error": str(e)}

backend/app/services/claude_service.py

+13-24
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,9 @@ class ClaudeService:
88
def __init__(self):
99
self.default_client = Anthropic()
1010

11-
def call_claude_api(self, system_prompt: str, data: dict, api_key: str | None = None) -> str:
11+
def call_claude_api(
12+
self, system_prompt: str, data: dict, api_key: str | None = None
13+
) -> str:
1214
"""
1315
Makes an API call to Claude and returns the response.
1416
@@ -32,40 +34,30 @@ def call_claude_api(self, system_prompt: str, data: dict, api_key: str | None =
3234
temperature=0,
3335
system=system_prompt,
3436
messages=[
35-
{
36-
"role": "user",
37-
"content": [
38-
{
39-
"type": "text",
40-
"text": user_message
41-
}
42-
]
43-
}
44-
]
37+
{"role": "user", "content": [{"type": "text", "text": user_message}]}
38+
],
4539
)
4640
return message.content[0].text # type: ignore
4741

48-
# autopep8: off
4942
def _format_user_message(self, data: dict[str, str]) -> str:
5043
"""Helper method to format the data into a user message"""
5144
parts = []
5245
for key, value in data.items():
53-
if key == 'file_tree':
46+
if key == "file_tree":
5447
parts.append(f"<file_tree>\n{value}\n</file_tree>")
55-
elif key == 'readme':
48+
elif key == "readme":
5649
parts.append(f"<readme>\n{value}\n</readme>")
57-
elif key == 'explanation':
50+
elif key == "explanation":
5851
parts.append(f"<explanation>\n{value}\n</explanation>")
59-
elif key == 'component_mapping':
52+
elif key == "component_mapping":
6053
parts.append(f"<component_mapping>\n{value}\n</component_mapping>")
61-
elif key == 'instructions' and value != "":
54+
elif key == "instructions" and value != "":
6255
parts.append(f"<instructions>\n{value}\n</instructions>")
63-
elif key == 'diagram':
56+
elif key == "diagram":
6457
parts.append(f"<diagram>\n{value}\n</diagram>")
65-
elif key == 'explanation':
58+
elif key == "explanation":
6659
parts.append(f"<explanation>\n{value}\n</explanation>")
6760
return "\n\n".join(parts)
68-
# autopep8: on
6961

7062
def count_tokens(self, prompt: str) -> int:
7163
"""
@@ -79,9 +71,6 @@ def count_tokens(self, prompt: str) -> int:
7971
"""
8072
response = self.default_client.messages.count_tokens(
8173
model="claude-3-5-sonnet-latest",
82-
messages=[{
83-
"role": "user",
84-
"content": prompt
85-
}]
74+
messages=[{"role": "user", "content": prompt}],
8675
)
8776
return response.input_tokens

0 commit comments

Comments
 (0)