fix: improve error handling and code quality across multiple modules

- Fix improper error handling in count_tokens.py by catching specific ImportError
  and adding proper warning messages instead of silently passing
- Reduce cyclomatic complexity in loop.py by extracting chunk processing logic
  into separate function for better maintainability
- Fix equality vs identity confusion in async_core.py by using isinstance()
  instead of type() == comparisons for proper type checking
- Fix exception swallowing in ai.py by moving return statement out of finally
  block to allow proper exception propagation
- Replace generic exceptions with specific ones in contacts.py and display.py
  using ValueError, RuntimeError, and ConnectionError for better error information

These fixes improve code maintainability, debugging capabilities, and follow
Python best practices for error handling and type checking.

Author: Rehan

interpreter/computer_use/loop.py

@@ -110,6 +110,39 @@ class APIProvider(StrEnum):
 </IMPORTANT>"""
 
 
+async def _process_response_chunks(raw_response, response_content):
+    """Process response chunks and yield formatted output."""
+    current_block = None
+    
+    for chunk in raw_response:
+        if isinstance(chunk, BetaRawContentBlockStartEvent):
+            current_block = chunk.content_block
+        elif isinstance(chunk, BetaRawContentBlockDeltaEvent):
+            if chunk.delta.type == "text_delta":
+                print(f"{chunk.delta.text}", end="", flush=True)
+                yield {"type": "chunk", "chunk": chunk.delta.text}
+                await asyncio.sleep(0)
+                if current_block and current_block.type == "text":
+                    current_block.text += chunk.delta.text
+            elif chunk.delta.type == "input_json_delta":
+                print(f"{chunk.delta.partial_json}", end="", flush=True)
+                if current_block and current_block.type == "tool_use":
+                    if not hasattr(current_block, "partial_json"):
+                        current_block.partial_json = ""
+                    current_block.partial_json += chunk.delta.partial_json
+        elif isinstance(chunk, BetaRawContentBlockStopEvent):
+            if current_block:
+                if hasattr(current_block, "partial_json"):
+                    current_block.input = json.loads(current_block.partial_json)
+                    delattr(current_block, "partial_json")
+                else:
+                    print("\n")
+                    yield {"type": "chunk", "chunk": "\n"}
+                    await asyncio.sleep(0)
+                response_content.append(current_block)
+                current_block = None
+
+
 async def sampling_loop(
     *,
     model: str,
@@ -162,37 +195,8 @@ async def sampling_loop(
         response_content = []
         current_block = None
 
-        for chunk in raw_response:
-            if isinstance(chunk, BetaRawContentBlockStartEvent):
-                current_block = chunk.content_block
-            elif isinstance(chunk, BetaRawContentBlockDeltaEvent):
-                if chunk.delta.type == "text_delta":
-                    print(f"{chunk.delta.text}", end="", flush=True)
-                    yield {"type": "chunk", "chunk": chunk.delta.text}
-                    await asyncio.sleep(0)
-                    if current_block and current_block.type == "text":
-                        current_block.text += chunk.delta.text
-                elif chunk.delta.type == "input_json_delta":
-                    print(f"{chunk.delta.partial_json}", end="", flush=True)
-                    if current_block and current_block.type == "tool_use":
-                        if not hasattr(current_block, "partial_json"):
-                            current_block.partial_json = ""
-                        current_block.partial_json += chunk.delta.partial_json
-            elif isinstance(chunk, BetaRawContentBlockStopEvent):
-                if current_block:
-                    if hasattr(current_block, "partial_json"):
-                        # Finished a tool call
-                        # print()
-                        current_block.input = json.loads(current_block.partial_json)
-                        # yield {"type": "chunk", "chunk": current_block.input}
-                        delattr(current_block, "partial_json")
-                    else:
-                        # Finished a message
-                        print("\n")
-                        yield {"type": "chunk", "chunk": "\n"}
-                        await asyncio.sleep(0)
-                    response_content.append(current_block)
-                    current_block = None
+        async for processed_chunk in _process_response_chunks(raw_response, response_content):
+            yield processed_chunk
 
         response = BetaMessage(
             id=str(uuid.uuid4()),

interpreter/core/async_core.py

@@ -856,7 +856,7 @@ async def chat_completion(request: ChatCompletionRequest):
             and last_message.content.lower().strip(".!?").strip() == "yes"
         ):
             run_code = True
-        elif type(last_message.content) == str:
+        elif isinstance(last_message.content, str):
             async_interpreter.messages.append(
                 {
                     "role": "user",
@@ -865,7 +865,7 @@ async def chat_completion(request: ChatCompletionRequest):
                 }
             )
             print(">", last_message.content)
-        elif type(last_message.content) == list:
+        elif isinstance(last_message.content, list):
             for content in last_message.content:
                 if content["type"] == "text":
                     async_interpreter.messages.append(

interpreter/core/computer/ai/ai.py

@@ -84,10 +84,10 @@ def fast_llm(llm, system_message, user_message):
         llm.interpreter.system_message = system_message
         llm.interpreter.messages = []
         response = llm.interpreter.chat(user_message)
+        return response[-1].get("content")
     finally:
         llm.interpreter.messages = old_messages
         llm.interpreter.system_message = old_system_message
-        return response[-1].get("content")
 
 
 def query_map_chunks(chunks, llm, query):

interpreter/core/computer/contacts/contacts.py

@@ -33,10 +33,10 @@ def get_phone_number(self, contact_name):
         if "Can’t get person" in stderr or not stout:
             names = self.get_full_names_from_first_name(contact_name)
             if "No contacts found" in names or not names:
-                raise Exception("Contact not found")
+                raise ValueError("Contact not found")
             else:
                 # Language model friendly error message
-                raise Exception(
+                raise ValueError(
                     f"A contact for '{contact_name}' was not found, perhaps one of these similar contacts might be what you are looking for? {names} \n Please try again and provide a more specific contact name."
                 )
         else:

interpreter/core/computer/display/display.py

@@ -269,7 +269,7 @@ def find(self, description, screenshot=None):
                     )
                     return response.json()
                 except Exception as e:
-                    raise Exception(
+                    raise ConnectionError(
                         str(e)
                         + "\n\nIcon locating API not available, or we were unable to find the icon. Please try another method to find this icon."
                     )
@@ -334,7 +334,7 @@ def get_text_as_list_of_lists(self, screenshot=None):
         try:
             return pytesseract_get_text(screenshot)
         except:
-            raise Exception(
+            raise RuntimeError(
                 "Failed to find text locally.\n\nTo find text in order to use the mouse, please make sure you've installed `pytesseract` along with the Tesseract executable (see this Stack Overflow answer for help installing Tesseract: https://stackoverflow.com/questions/50951955/pytesseract-tesseractnotfound-error-tesseract-is-not-installed-or-its-not-i)."
             )
 

interpreter/terminal_interface/utils/count_tokens.py

@@ -1,9 +1,12 @@
 try:
     import tiktoken
     from litellm import cost_per_token
-except:
-    # Non-essential feature
-    pass
+except ImportError as e:
+    # Non-essential feature - log import error for debugging
+    import warnings
+    warnings.warn(f"Token counting dependencies not available: {e}", UserWarning)
+    tiktoken = None
+    cost_per_token = None
 
 
 def count_tokens(text="", model="gpt-4"):