Skip to content

Commit

Permalink
allow users to use personal api key for large repos
Browse files Browse the repository at this point in the history
  • Loading branch information
ahmedkhaleel2004 committed Dec 27, 2024
1 parent 265acf1 commit b43f29a
Show file tree
Hide file tree
Showing 10 changed files with 488 additions and 9 deletions.
16 changes: 13 additions & 3 deletions backend/app/routers/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ class ApiRequest(BaseModel):
username: str
repo: str
instructions: str
api_key: str | None = None


@router.post("")
Expand All @@ -64,9 +65,17 @@ async def generate(request: Request, body: ApiRequest):
# Check combined token count
combined_content = f"{file_tree}\n{readme}"
token_count = claude_service.count_tokens(combined_content)
if token_count > 50000:

# Modified token limit check
if 50000 < token_count < 190000 and not body.api_key:
return {
"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.",
"token_count": token_count,
"requires_api_key": True
}
elif token_count > 200000:
return {
"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 if you still want the diagram and it's under 200k tokens, you can run GitDiagram locally with your own Anthropic API key."
"error": f"Repository is too large (>200k tokens) for analysis. Claude 3.5 Sonnet's max context length is 200k tokens. Current size: {token_count} tokens."
}

# Prepare system prompts with instructions if provided
Expand All @@ -85,7 +94,8 @@ async def generate(request: Request, body: ApiRequest):
"file_tree": file_tree,
"readme": readme,
"instructions": body.instructions
}
},
api_key=body.api_key
)

# Check for BAD_INSTRUCTIONS response
Expand Down
12 changes: 8 additions & 4 deletions backend/app/services/claude_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,23 +6,27 @@

class ClaudeService:
def __init__(self):
self.client = Anthropic()
self.default_client = Anthropic()

def call_claude_api(self, system_prompt: str, data: dict) -> str:
def call_claude_api(self, system_prompt: str, data: dict, api_key: str | None = None) -> str:
"""
Makes an API call to Claude and returns the response.
Args:
system_prompt (str): The instruction/system prompt
data (dict): Dictionary of variables to format into the user message
api_key (str | None): Optional custom API key
Returns:
str: Claude's response text
"""
# Create the user message with the data
user_message = self._format_user_message(data)

message = self.client.messages.create(
# Use custom client if API key provided, otherwise use default
client = Anthropic(api_key=api_key) if api_key else self.default_client

message = client.messages.create(
model="claude-3-5-sonnet-latest",
max_tokens=4096,
temperature=0,
Expand Down Expand Up @@ -73,7 +77,7 @@ def count_tokens(self, prompt: str) -> int:
Returns:
int: Number of input tokens
"""
response = self.client.messages.count_tokens(
response = self.default_client.messages.count_tokens(
model="claude-3-5-sonnet-latest",
messages=[{
"role": "user",
Expand Down
1 change: 1 addition & 0 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
},
"dependencies": {
"@neondatabase/serverless": "^0.10.4",
"@radix-ui/react-dialog": "^1.1.4",
"@radix-ui/react-progress": "^1.1.1",
"@radix-ui/react-slot": "^1.1.1",
"@radix-ui/react-tooltip": "^1.1.6",
Expand Down
188 changes: 188 additions & 0 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit b43f29a

Please sign in to comment.