1515import requests
1616import openai
1717from anthropic import Anthropic
18+ import tiktoken
1819
1920
2021class OllamaAPI :
2122 def __init__ (self , model ):
2223 self .model = model
2324 self .base_url = "http://localhost:11434/api"
25+ self .token_count = 0
2426
27+ def count_tokens (self , text ):
28+ # Ollama doesn't provide a built-in token counter, so we'll use tiktoken as an approximation
29+ return len (tiktoken .encoding_for_model ("gpt-4o" ).encode (text ))
30+
2531 def generate (self , prompt ):
2632 url = f"{ self .base_url }
2733 data = {"model" : self .model , "prompt" : prompt , "stream" : True }
@@ -39,6 +45,18 @@ def generate(self, prompt):
3945 except json .JSONDecodeError :
4046 print (f"Error decoding JSON: { decoded_line } )
4147 print () # Print a newline at the end
48+
49+ # Extract content between markers
50+ start_marker = "^^^start^^^"
51+ end_marker = "^^^end^^^"
52+ start_index = full_response .find (start_marker )
53+ end_index = full_response .find (end_marker )
54+ if start_index != - 1 and end_index != - 1 :
55+ full_response = full_response [start_index + len (start_marker ):end_index ].strip ()
56+
57+ self .token_count = self .count_tokens (full_response )
58+ print (f"Token count: { self .token_count } )
59+
4260 return full_response
4361 else :
4462 raise Exception (f"Ollama API error: { response .text } )
@@ -53,6 +71,10 @@ def __init__(self, model):
5371 if not self .api_key :
5472 raise ValueError ("OPENAI_API_KEY environment variable is not set" )
5573 openai .api_key = self .api_key
74+ self .token_count = 0
75+
76+ def count_tokens (self , text ):
77+ return len (tiktoken .encoding_for_model (self .model ).encode (text ))
5678
5779 def generate (self , prompt ):
5880 try :
@@ -68,11 +90,24 @@ def generate(self, prompt):
6890 full_response += chunk_text
6991 print (chunk_text , end = "" , flush = True )
7092 print () # Print a newline at the end
93+
94+ # Extract content between markers
95+ start_marker = "^^^start^^^"
96+ end_marker = "^^^end^^^"
97+ start_index = full_response .find (start_marker )
98+ end_index = full_response .find (end_marker )
99+ if start_index != - 1 and end_index != - 1 :
100+ full_response = full_response [start_index + len (start_marker ):end_index ].strip ()
101+
102+ self .token_count = self .count_tokens (full_response )
103+ print (f"Token count: { self .token_count } )
104+
71105 return full_response
72106 except Exception as e :
73107 raise Exception (f"OpenAI API error: { str (e )} )
74108
75109
110+
76111class ClaudeAPI :
77112 def __init__ (self , model ):
78113 if model == "mistral-nemo" :
@@ -82,22 +117,53 @@ def __init__(self, model):
82117 if not self .api_key :
83118 raise ValueError ("ANTHROPIC_API_KEY environment variable is not set" )
84119 self .client = Anthropic (api_key = self .api_key )
120+ self .token_count = 0
121+
122+ def count_tokens (self , text ):
123+ # Ollama doesn't provide a built-in token counter, so we'll use tiktoken as an approximation
124+ return len (tiktoken .encoding_for_model ("gpt-4o" ).encode (text ))
85125
86126 def generate (self , prompt ):
87127 try :
88- response = self .client .messages .create (
89- model = self .model ,
90- messages = [{"role" : "user" , "content" : prompt }],
91- stream = True ,
92- max_tokens = 1000 ,
93- )
94128 full_response = ""
95- for completion in response :
96- if completion .type == "content_block_delta" :
97- chunk_text = completion .delta .text
98- full_response += chunk_text
99- print (chunk_text , end = "" , flush = True )
129+ max_iterations = 5 # Adjust this value as needed
130+ continuation_prompt = prompt
131+
132+ for iteration in range (max_iterations ):
133+ response = self .client .messages .create (
134+ model = self .model ,
135+ messages = [{"role" : "user" , "content" : continuation_prompt }],
136+ stream = True ,
137+ max_tokens = 1000 ,
138+ )
139+
140+ chunk_response = ""
141+ for completion in response :
142+ if completion .type == "content_block_delta" :
143+ chunk_text = completion .delta .text
144+ chunk_response += chunk_text
145+ print (chunk_text , end = "" , flush = True )
146+
147+ full_response += chunk_response
148+
149+ if "^^^end^^^" in chunk_response :
150+ break
151+
152+ continuation_prompt = f"Continue from where you left off. Previous response: { chunk_response }
153+
100154 print () # Print a newline at the end
155+
156+ # Extract content between markers
157+ start_marker = "^^^start^^^"
158+ end_marker = "^^^end^^^"
159+ start_index = full_response .find (start_marker )
160+ end_index = full_response .find (end_marker )
161+ if start_index != - 1 and end_index != - 1 :
162+ full_response = full_response [start_index + len (start_marker ):end_index ].strip ()
163+
164+ self .token_count = self .count_tokens (full_response )
165+ print (f"Token count: { self .token_count } )
166+
101167 return full_response
102168 except Exception as e :
103169 raise Exception (f"Claude API error: { str (e )} )
@@ -119,6 +185,11 @@ def __init__(
119185 self .pwd = os .getcwd () + "/" + self .project_name
120186 self .llm = self .setup_llm ()
121187 self .previous_suggestions = set ()
188+ self .token_counts = {}
189+
190+ def count_tokens (self , text ):
191+ encoding = tiktoken .encoding_for_model ("gpt-3.5-turbo" )
192+ return len (encoding .encode (text ))
122193
123194 def setup_llm (self ):
124195 if self .provider == "ollama" :
@@ -169,6 +240,9 @@ def run_task(self):
169240 break
170241 test_check_attempts += 1
171242
243+ total_tokens = sum (self .token_counts .values ())
244+ print (f"\n Total tokens used: { total_tokens } )
245+
172246 print (
173247 "Task completed. Please review the output and make any necessary manual adjustments."
174248 )
@@ -303,7 +377,7 @@ def implement_solution(self, max_attempts=3):
303377 prompt = f"""
304378 Create a comprehensive implementation for the task: { self .task }
305379 You must follow these rules strictly:
306- 1, IMPORTANT: Never use pass statements in your code or tests. Always provide a meaningful implementation.
380+ 1. IMPORTANT: Never use pass statements in your code or tests. Always provide a meaningful implementation.
307381 2. CRITICAL: Use the following code block format for specifying file content:
308382 For code files, use:
309383 <<<main.py>>>
@@ -332,6 +406,7 @@ def implement_solution(self, max_attempts=3):
332406 13. IMPORTANT: Always pytest parameterize tests for different cases.
333407 14. CRITICAL: Always use `import main` to import the main.py file in the test file.
334408 15. IMPORTANT: Only mock external services or APIs in tests.
409+ 16. IMPORTANT: Enclose your entire response between ^^^start^^^ and ^^^end^^^ markers.
335410 Working directory: { self .pwd }
336411 """
337412
@@ -340,6 +415,14 @@ def implement_solution(self, max_attempts=3):
340415 solution = self .get_response (prompt )
341416 self .logger .info (f"Received solution:\n { solution } )
342417
418+ # Extract content between markers
419+ start_marker = "^^^start^^^"
420+ end_marker = "^^^end^^^"
421+ start_index = solution .find (start_marker )
422+ end_index = solution .find (end_marker )
423+ if start_index != - 1 and end_index != - 1 :
424+ solution = solution [start_index + len (start_marker ):end_index ].strip ()
425+
343426 # Parse and execute any uv add commands
344427 uv_commands = [
345428 line .strip ()
@@ -392,11 +475,14 @@ def extract_file_contents_direct(self, solution):
392475
393476 def get_response (self , prompt ):
394477 try :
395- return self .llm .generate (prompt )
478+ response = self .llm .generate (prompt )
479+ prompt_key = prompt [:50 ] # Use first 50 characters as a key
480+ self .token_counts [prompt_key ] = self .llm .token_count
481+ return response
396482 except Exception as e :
397483 self .logger .error (f"Error getting response from { self .provider } { str (e )} )
398484 return ""
399-
485+
400486 def code_check (self , file_path ):
401487 try :
402488 # Run autopep8 to automatically fix style issues
@@ -498,6 +584,7 @@ def improve_test_file(self, test_output):
498584 8. IMPORTANT: Always pytest parameterize tests for different cases.
499585 9. CRITICAL: Always use `import main` to import the main.py file in the test file.
500586 10. IMPORTANT: Only mock external services or APIs in tests.
587+ 11. IMPORTANT: Enclose your entire response between ^^^start^^^ and ^^^end^^^ markers.
501588 Working directory: { self .pwd }
502589 """
503590 proposed_improvements = self .get_response (prompt )
@@ -522,6 +609,7 @@ def validate_implementation(self, proposed_improvements):
522609 If the implementation is correct or mostly correct, respond with 'VALID'.
523610 If the implementation is completely unrelated or fundamentally flawed, respond with 'INVALID'.
524611 Do not provide any additional information or explanations.
612+ IMPORTANT: Enclose your entire response between ^^^start^^^ and ^^^end^^^ markers.
525613 """
526614 response = self .get_response (prompt )
527615
@@ -567,6 +655,7 @@ def improve_code(
567655 # File content here
568656 <<<end>>>
569657 7. CRITICAL: Do not explain the task only implement the required functionality in the code blocks.
658+ 8. IMPORTANT: Enclose your entire response between ^^^start^^^ and ^^^end^^^ markers.
570659 Working directory: { self .pwd }
571660 """
572661 proposed_improvements = self .get_response (prompt )
0 commit comments