Skip to content

Commit

Permalink
fix aider.py
Browse files Browse the repository at this point in the history
  • Loading branch information
waleko committed Jul 10, 2024
1 parent 052f826 commit 33e6676
Showing 1 changed file with 14 additions and 1 deletion.
15 changes: 14 additions & 1 deletion code_editing/agents/context_providers/aider/aider.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import tiktoken

from code_editing.agents.context_providers.aider.repo_map import RepoMap, find_src_files
from code_editing.agents.context_providers.context_provider import ContextProvider

Expand All @@ -7,10 +9,21 @@ def __init__(self, repo_path: str, data_path: str):
self.repo_path = repo_path
self.data_path = data_path

def count_tokens(text, model="gpt-4o"):
# Get the appropriate encoding for the specified model
encoding = tiktoken.encoding_for_model(model)

# Encode the text using the encoding scheme
tokens = encoding.encode(text)

# Return the number of tokens
return len(tokens)

self.rm = RepoMap(
map_tokens=1024,
root=repo_path,
token_count=lambda x: len(x.split()) // 4,
token_count=count_tokens,
# max_context_window=120000
)

def get_repo_map(self) -> str:
Expand Down

0 comments on commit 33e6676

Please sign in to comment.