1  #!/usr/bin/env python3
  2  """
  3  Simple script to add session_consult tool to all ECHO agents.
  4  """
  5  
  6  import re
  7  import os
  8  import subprocess
  9  
 10  AGENTS = {
 11      "cto": "cto",
 12      "chro": "chro",
 13      "operations_head": "operations_head",
 14      "product_manager": "product_manager",
 15      "senior_architect": "senior_architect",
 16      "uiux_engineer": "uiux_engineer",
 17      "senior_developer": "senior_developer",
 18      "test_lead": "test_lead",
 19  }
 20  
 21  TOOL_DEF_TEMPLATE = '''      },
 22        %{
 23          name: "session_consult",
 24          description: """
 25          Query the AI assistant with conversation memory (LocalCode-style).
 26  
 27          Maintains multi-turn conversations with automatic context injection:
 28          - Your role, responsibilities, and authority limits
 29          - Recent decisions and messages (last 5 each)
 30          - Current system status (PostgreSQL, Redis, Ollama)
 31          - Git context (branch, last commit)
 32          - Conversation history (last 5 turns)
 33  
 34          Perfect for exploratory questions, decision analysis with iterative thinking,
 35          and strategy planning with follow-up questions.
 36          """,
 37          inputSchema: %{
 38            type: "object",
 39            properties: %{
 40              question: %{
 41                type: "string",
 42                description: "The question to ask the AI assistant",
 43                minLength: 1
 44              },
 45              session_id: %{
 46                type: "string",
 47                description: "Session ID to continue conversation (optional, omit for new session)"
 48              },
 49              context: %{
 50                type: "string",
 51                description: "Additional context for this specific query (optional)"
 52              }
 53            },
 54            required: ["question"]
 55          }
 56        }'''
 57  
 58  def add_session_consult(agent_dir, role):
 59      """Add session_consult tool to an agent."""
 60      file_path = f"apps/{agent_dir}/lib/{agent_dir}.ex"
 61  
 62      if not os.path.exists(file_path):
 63          print(f"⚠️  Skipping {agent_dir}: File not found")
 64          return False
 65  
 66      with open(file_path, 'r') as f:
 67          content = f.read()
 68  
 69      # Check if already has session_consult
 70      if '"session_consult"' in content:
 71          print(f"   ✓ {agent_dir} already has session_consult")
 72          return True
 73  
 74      # Backup
 75      with open(f"{file_path}.backup", 'w') as f:
 76          f.write(content)
 77  
 78      #Step 1: Add tool definition before closing ]
 79      # Find the last tool in the list (look for pattern "}\n    ]\n  end" in tools function)
 80      tool_pattern = r'(\s+})\s+(\]\s+end\s+@impl true\s+def execute_tool)'
 81  
 82      if not re.search(tool_pattern, content):
 83          print(f"   ❌ {agent_dir}: Could not find tools list end")
 84          return False
 85  
 86      content = re.sub(
 87          tool_pattern,
 88          rf'\1,{TOOL_DEF_TEMPLATE}\n    \2',
 89          content
 90      )
 91  
 92      # Step 2: Add execute_tool handler before catch-all handler
 93      handler_code = f'''
 94    def execute_tool("session_consult", args) do
 95      question = Map.fetch!(args, "question")
 96      session_id = Map.get(args, "session_id")
 97      context = Map.get(args, "context")
 98  
 99      opts = if context, do: [context: context], else: []
100  
101      case DecisionHelper.consult_session(:{role}, session_id, question, opts) do
102        {{:ok, result}} ->
103          response = format_session_response(result)
104          {{:ok, response}}
105  
106        {{:error, :llm_disabled}} ->
107          {{:error, "LLM is disabled for {agent_dir}. Enable with LLM_ENABLED=true or {agent_dir.upper()}_LLM_ENABLED=true"}}
108  
109        {{:error, :session_not_found}} ->
110          {{:error, "Session not found: #{{session_id}}. It may have expired after 1 hour of inactivity."}}
111  
112        {{:error, reason}} ->
113          {{:error, "AI consultation failed: #{{inspect(reason)}}"}}
114      end
115    end
116  '''
117  
118      # Find catch-all handler: def execute_tool(name, _args) do
119      catchall_pattern = r'(\s+def execute_tool\(name, _args\) do)'
120      content = re.sub(catchall_pattern, rf'{handler_code}\1', content)
121  
122      # Step 3: Add format_session_response helper before module end
123      helper_code = f'''
124    defp format_session_response(result) do
125      model = EchoShared.LLM.Config.get_model(:{role})
126  
127      base = %{{
128        "response" => result.response,
129        "session_id" => result.session_id,
130        "turn_count" => result.turn_count,
131        "estimated_tokens" => result.total_tokens,
132        "model" => model,
133        "agent" => "{agent_dir}"
134      }}
135  
136      if result.warnings != [] do
137        Map.put(base, "warnings", result.warnings)
138      else
139        base
140      end
141    end
142  '''
143  
144      # Add before final "end"
145      content = content.rstrip() + helper_code + "\nend\n"
146  
147      # Write modified content
148      with open(file_path, 'w') as f:
149          f.write(content)
150  
151      # Compile to check
152      result = subprocess.run(
153          ['mix', 'compile'],
154          cwd=f"apps/{agent_dir}",
155          capture_output=True,
156          text=True
157      )
158  
159      if result.returncode == 0 and 'Generated' in result.stdout:
160          print(f"   ✅ {agent_dir}: Success!")
161          os.remove(f"{file_path}.backup")
162          return True
163      else:
164          print(f"   ❌ {agent_dir}: Compilation failed!")
165          print(f"      {result.stderr[:200]}")
166          # Restore backup
167          with open(f"{file_path}.backup", 'r') as f:
168              with open(file_path, 'w') as out:
169                  out.write(f.read())
170          return False
171  
172  def main():
173      print("=" * 50)
174      print(" Adding session_consult to all agents")
175      print("=" * 50)
176      print()
177  
178      success_count = 0
179      for agent_dir, role in AGENTS.items():
180          print(f"📝 Processing {agent_dir} (role: {role})...")
181          if add_session_consult(agent_dir, role):
182              success_count += 1
183          print()
184  
185      print("=" * 50)
186      print(f" Completed: {success_count}/{len(AGENTS)} agents updated")
187      print("=" * 50)
188  
189  if __name__ == "__main__":
190      main()