code comibner fixes

app.py

@@ -474,12 +474,12 @@ async def chat_with_all(
                             "status": "error"
                         }) + "\n"
                         return
-                    if agent_name != "code_combiner_agent":
+                    if "code_combiner_agent" in agent_name:
                         total_response += str(response) if response else ""
                         total_inputs += str(inputs) if inputs else ""
 
                     yield json.dumps({
-                        "agent": agent_name,
+                        "agent": agent_name.split("__")[0] if "__" in agent_name else agent_name,
                         "content": formatted_response,
                         "status": "success" if response else "error"
                     }) + "\n"
@@ -490,7 +490,7 @@ async def chat_with_all(
 
                 # Track the code combiner agent response
                 if "refined_complete_code" in response:
-                    model_name = "claude-3-5-sonnet-latest"
+                    model_name = agent_name.split("__")[0] if "__" in agent_name else agent_name
                     provider = app.state.ai_manager.get_provider_for_model(model_name)
                     input_tokens = len(app.state.ai_manager.tokenizer.encode(str(inputs)))
                     completion_tokens = len(app.state.ai_manager.tokenizer.encode(str(response)))

scripts/format_response.py

@@ -28,6 +28,10 @@ def clean_print_statements(code_block):
     # This regex targets print statements, even if they have newlines inside
     return re.sub(r'print\((.*?)(\\n.*?)(.*?)\)', r'print(\1\3)', code_block, flags=re.DOTALL)
 
+def remove_code_block_from_summary(summary):
+    # use regex to remove code block from summary list
+    summary = re.sub(r'```python\n(.*?)\n```', '', summary)
+    return summary.split("\n")
 
 def remove_main_block(code):
     # Match the __main__ block
@@ -130,6 +134,9 @@ def format_code_backticked_block(code_str):
     # Remove empty lines that might have been created after cleaning
     code_clean = re.sub(r"\n\s*\n+", "\n\n", code_clean)
     
+    # remove main
+    code_clean = remove_main_block(code_clean)
+    
     return f'```python\n{code_clean}\n```'
 
     
@@ -267,16 +274,17 @@ def execute_code_from_markdown(code_str, dataframe=None):
 def format_response_to_markdown(api_response, agent_name = None, dataframe=None):
     try:
         markdown = []
-                
-        # Handle error responses
-        if isinstance(api_response, dict) and "error" in api_response:
-            return f"**Error**: {api_response['error']}"
+        logger.log_message(f"Formatting response to markdown for agent '{agent_name}' at {time.strftime('%Y-%m-%d %H:%M:%S')}: {api_response}, length: {len(api_response)}", level=logging.INFO)
         
         if isinstance(api_response, dict):
             for key in api_response:
                 if "error" in api_response[key]:
                     return f"**Error**: Rate limit exceeded. Please try switching models from the settings."
                 # You can add more checks here if needed for other keys
+                       
+        # Handle error responses
+        if isinstance(api_response, dict) and "error" in api_response:
+            return f"**Error**: {api_response['error']}"
         if "response" in api_response and isinstance(api_response['response'], str):
             if any(err in api_response['response'].lower() for err in ["auth", "api", "lm"]):
                 return "**Error**: Authentication failed. Please check your API key in settings and try again."
@@ -284,6 +292,7 @@ def format_response_to_markdown(api_response, agent_name = None, dataframe=None)
                 return "**Error**: Model configuration error. Please verify your model selection in settings."
 
         for agent, content in api_response.items():
+            agent = agent.split("__")[0] if "__" in agent else agent
             if "memory" in agent or not content:
                 continue
                 
@@ -315,21 +324,30 @@ def format_response_to_markdown(api_response, agent_name = None, dataframe=None)
 
             if 'summary' in content:
                 # make the summary a bullet-point list
+                summary_lines = remove_code_block_from_summary(content['summary'])
                 summary_lines = content['summary'].split('\n')
+                # remove code block from summary
                 markdown.append("### Summary\n")
                 for line in summary_lines:
                     markdown.append(f"• {line.strip().replace('•', '').replace('-', '')}\n")
 
-            if 'refined_complete_code' in content:
+            if 'refined_complete_code' in content and 'summary' in content:
+                logger.log_message(f"Refined complete code: {content['refined_complete_code']}", level=logging.INFO)
                 try:
-                    clean_code = format_code_block(content['refined_complete_code']) 
-                    output, json_outputs = execute_code_from_markdown(clean_code, dataframe)
+                    if content['refined_complete_code'] is not None and content['refined_complete_code'] != "":
+                        clean_code = format_code_block(content['refined_complete_code']) 
+                        markdown_code = format_code_backticked_block(content['refined_complete_code'])
+                        output, json_outputs = execute_code_from_markdown(clean_code, dataframe)
+                    elif "```python" in content['summary']:
+                        clean_code = format_code_block(content['summary'])
+                        markdown_code = format_code_backticked_block(content['summary'])
+                        output, json_outputs = execute_code_from_markdown(clean_code, dataframe)
                 except Exception as e:
                     logger.log_message(f"Error in execute_code_from_markdown: {str(e)}", level=logging.ERROR)
                     markdown.append(f"**Error**: {str(e)}")
                     # continue
                 
-                markdown.append(f"### Refined Complete Code\n{format_code_backticked_block(content['refined_complete_code'])}\n")
+                markdown.append(f"### Refined Complete Code\n{markdown_code}\n")
 
                 if output:
                     markdown.append("### Execution Output\n")
@@ -339,7 +357,7 @@ def format_response_to_markdown(api_response, agent_name = None, dataframe=None)
                     markdown.append("### Plotly JSON Outputs\n")
                     for idx, json_output in enumerate(json_outputs):
                         markdown.append(f"```plotly\n{json_output}\n```\n")
-
+                logger.log_message(f"Refined complete code executed, {markdown}", level=logging.INFO)
             # if agent_name is not None:  
             #     if f"memory_{agent_name}" in api_response:
             #         markdown.append(f"### Memory\n{api_response[f'memory_{agent_name}']}\n")
@@ -347,7 +365,8 @@ def format_response_to_markdown(api_response, agent_name = None, dataframe=None)
     except Exception as e:
         logger.log_message(f"Error in format_response_to_markdown: {str(e)}", level=logging.ERROR)
         return f"{str(e)}"
-    logger.log_message(f"Generated markdown content for agent '{agent_name}' at {time.strftime('%Y-%m-%d %H:%M:%S')}: {api_response}, length: {len(markdown)}", level=logging.INFO)
+    
+    logger.log_message(f"Generated markdown content for agent '{agent_name}' at {time.strftime('%Y-%m-%d %H:%M:%S')}: {markdown}, length: {len(markdown)}", level=logging.INFO)
     
     
     if not markdown or len(markdown) <= 1:

scripts/tier_maker.py

@@ -37,7 +37,8 @@ costs = {
         "gemma-7b-it": {"input": 0.00007, "output": 0.00007},
         "gemma2-9b-it": {"input": 0.0002, "output": 0.0002},
         "llama3-groq-70b-8192-tool-use-preview": {"input": 0.00089, "output": 0.00089},
-        "llama3-groq-8b-8192-tool-use-preview": {"input": 0.00019, "output": 0.00019}
+        "llama3-groq-8b-8192-tool-use-preview": {"input": 0.00019, "output": 0.00019},
+        "qwen-2.5-coder-32b": {"input": 0.0015, "output": 0.003}
     }
 }
         
@@ -139,6 +140,7 @@ print(json.dumps(model_tiers, indent=4))
             "claude-3-opus-latest",
             "claude-3-7-sonnet-latest",
             "claude-3-5-sonnet-latest",
+            "qwen-2.5-coder-32b",
             "claude-3-5-haiku-latest",
             "deepseek-r1-distill-qwen-32b",
             "deepseek-r1-distill-llama-70b",

src/agents/agents.py

@@ -3,6 +3,9 @@ import src.agents.memory_agents as m
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import os
+from dotenv import load_dotenv
+
+load_dotenv()
 
 class analytical_planner(dspy.Signature):
     # The planner agent which routes the query to Agent(s)
@@ -461,26 +464,7 @@ class auto_analyst(dspy.Module):
         for i, a in enumerate(agents):
             name = a.__pydantic_core_schema__['schema']['model_name']
             self.agents[name] = dspy.ChainOfThought(a)
-            # parse agent inputs safely handling edge cases
-            agent_str = str(agents[i].__pydantic_core_schema__['cls'])
-            try:
-                # Handle cases when the string format might change or vary
-                if '(' in agent_str:
-                    # Extract the part between ( and the first -> ORR )
-                    params_part = agent_str.split('(', 1)[1]
-                    if '->' in params_part:
-                        params_part = params_part.split('->', 1)[0]
-                    if ')' in params_part:
-                        params_part = params_part.split(')', 1)[0]
-                    
-                    # Split by comma and clean up
-                    self.agent_inputs[name] = {x.strip() for x in params_part.split(',') if x.strip()}
-                else:
-                    # Default empty set if no parameters found
-                    self.agent_inputs[name] = set()
-            except (IndexError, Exception):
-                # Fallback to empty set if parsing fails
-                self.agent_inputs[name] = set()
+            self.agent_inputs[name] = {x.strip() for x in str(agents[i].__pydantic_core_schema__['cls']).split('->')[0].split('(')[1].split(',')}
             self.agent_desc.append(str(a.__pydantic_core_schema__['cls']))
         
         # Initialize coordination agents
@@ -551,21 +535,21 @@ class auto_analyst(dspy.Module):
         # Execute code combiner after all agents complete
         code_list = [result['code'] for _, result in completed_results if 'code' in result]
         try:
-            with dspy.settings.context(lm=dspy.LM(model="anthropic/claude-3-7-sonnet-latest", max_tokens=12000, temperature=1.0)):
+            with dspy.settings.context(lm=dspy.GROQ(model="qwen-2.5-coder-32b", max_tokens=12000, temperature=1.0, api_key=os.getenv("GROQ_API_KEY"))):
                 combiner_result = self.code_combiner_agent(agent_code_list=str(code_list), dataset=dict_['dataset'])
-                yield 'code_combiner_agent', str(code_list), dict(combiner_result)
+                yield 'code_combiner_agent__qwen', str(code_list), dict(combiner_result)
         except:
             try: 
-                with dspy.settings.context(lm=dspy.GROQ(model="deepseek-r1-distill-llama-70b", max_tokens=12000, temperature=1.0)):
+                with dspy.settings.context(lm=dspy.LM(model="anthropic/claude-3-7-sonnet-latest", max_tokens=12000, temperature=1.0)):
                     combiner_result = self.code_combiner_agent(agent_code_list=str(code_list), dataset=dict_['dataset'])
-                    yield 'code_combiner_agent', str(code_list), dict(combiner_result)
+                    yield 'code_combiner_agent__claude', str(code_list), dict(combiner_result)
             except:
                 try: 
-                    with dspy.settings.context(lm=dspy.GROQ(model="qwen-2.5-coder-32b", max_tokens=12000, temperature=1.0)):
+                    with dspy.settings.context(lm=dspy.GROQ(model="deepseek-r1-distill-llama-70b", max_tokens=12000, temperature=1.0, api_key=os.getenv("GROQ_API_KEY"))):
                         combiner_result = self.code_combiner_agent(agent_code_list=str(code_list), dataset=dict_['dataset'])
-                        yield 'code_combiner_agent', str(code_list), dict(combiner_result)
-                except:
-                    yield 'code_combiner_agent', str(code_list), {"error": "Error in code combiner"}
+                        yield 'code_combiner_agent__deepseek', str(code_list), dict(combiner_result)
+                except Exception as e:
+                    yield 'code_combiner_agent__none', str(code_list), {"error": "Error in code combiner: "+str(e)}
 
 # Agent to make a Chat history name from a query
 class chat_history_name_agent(dspy.Signature):

src/managers/ai_manager.py

@@ -52,7 +52,8 @@ costs = {
                 "gemma-7b-it": {"input": 0.00007, "output": 0.00007},
                 "gemma2-9b-it": {"input": 0.0002, "output": 0.0002},
                 "llama3-groq-70b-8192-tool-use-preview": {"input": 0.00089, "output": 0.00089},
-                "llama3-groq-8b-8192-tool-use-preview": {"input": 0.00019, "output": 0.00019}
+                "llama3-groq-8b-8192-tool-use-preview": {"input": 0.00019, "output": 0.00019},
+                "qwen-2.5-coder-32b": {"input": 0.0015, "output": 0.003}
             }
         }
         

src/managers/chat_manager.py

@@ -47,9 +47,29 @@ class ChatManager:
             "claude-3-haiku-20240307": {"input": 0.00025, "output": 0.00125},
             "claude-3-5-sonnet-latest": {"input": 0.003, "output": 0.015},
             # Groq models
-            "llama-3-70b-8192": {"input": 0.0007, "output": 0.0007},
-            "llama-3-8b-8192": {"input": 0.0002, "output": 0.0002},
-            "mixtral-8x7b-32768": {"input": 0.0006, "output": 0.0006},
+            "deepseek-r1-distill-qwen-32b": {"input": 0.00075, "output": 0.00099},
+            "deepseek-r1-distill-llama-70b": {"input": 0.00075, "output": 0.00099},
+            "llama-3.3-70b-versatile": {"input": 0.00059, "output": 0.00079},
+            "llama-3.3-70b-specdec": {"input": 0.00059, "output": 0.00099},
+            "llama2-70b-4096": {"input": 0.0007, "output": 0.0008},
+            "llama3-8b-8192": {"input": 0.00005, "output": 0.00008},
+            "llama-3.2-1b-preview": {"input": 0.00004, "output": 0.00004},
+            "llama-3.2-3b-preview": {"input": 0.00006, "output": 0.00006},
+            "llama-3.2-11b-text-preview": {"input": 0.00018, "output": 0.00018},
+            "llama-3.2-11b-vision-preview": {"input": 0.00018, "output": 0.00018},
+            "llama-3.2-90b-text-preview": {"input": 0.0009, "output": 0.0009},
+            "llama-3.2-90b-vision-preview": {"input": 0.0009, "output": 0.0009},
+            "llama3-70b-8192": {"input": 0.00059, "output": 0.00079},
+            "llama-3.1-8b-instant": {"input": 0.00005, "output": 0.00008},
+            "llama-3.1-70b-versatile": {"input": 0.00059, "output": 0.00079},
+            "llama-3.1-405b-reasoning": {"input": 0.00059, "output": 0.00079},
+            "mixtral-8x7b-32768": {"input": 0.00024, "output": 0.00024},
+            "gemma-7b-it": {"input": 0.00007, "output": 0.00007},
+            "gemma2-9b-it": {"input": 0.0002, "output": 0.0002},
+            "llama3-groq-70b-8192-tool-use-preview": {"input": 0.00089, "output": 0.00089},
+            "llama3-groq-8b-8192-tool-use-preview": {"input": 0.00019, "output": 0.00019},
+            "qwen-2.5-coder-32b": {"input": 0.0015, "output": 0.003}
+            
         }
         
         # Add model providers mapping