1  #!/usr/bin/env python3
  2  import os
  3  import sys
  4  import json
  5  import asyncio
  6  import aiohttp
  7  from datetime import datetime
  8  from typing import Dict, Any
  9  
 10  # ====================== CONFIGURATION ======================
 11  # ENDPOINT 1: OpenRouter (Cloud)
 12  OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
 13  OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
 14  OPENROUTER_MODEL = "openai/gpt-3.5-turbo"  # Free tier model
 15  
 16  # ENDPOINT 2: Groq (Cloud - Fast)
 17  GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
 18  GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"
 19  GROQ_MODEL = "llama-3.1-8b-instant"  # Fast free model
 20  
 21  # ENDPOINT 3: Local Ollama (Local)
 22  OLLAMA_URL = "http://localhost:11434/api/generate"
 23  OLLAMA_MODEL = "llama2:7b"  # Change to your local model
 24  
 25  # ENDPOINT 4: Direct Model API (Custom - e.g., io.net style)
 26  CUSTOM_API_URL = os.getenv("CUSTOM_API_URL", "http://localhost:8080/v1/chat/completions")
 27  CUSTOM_API_KEY = os.getenv("CUSTOM_API_KEY", "")
 28  CUSTOM_MODEL = "gpt-oss-20b"  # Or your deployed model
 29  
 30  # ====================== PROMPT ENGINEERING ======================
 31  def enhance_prompt(query: str) -> str:
 32      """Improve prompts for better responses"""
 33      query_lower = query.lower()
 34      
 35      # Math queries
 36      math_patterns = [r'\d+[\+\-\*\/]\d+', r'calculate', r'what is \d+', r'\d+\s*[\+\-\*\/]\s*\d+']
 37      import re
 38      for pattern in math_patterns:
 39          if re.search(pattern, query_lower):
 40              return f"Calculate this exactly: {query}. Give only the numerical answer without explanation."
 41      
 42      # Factual queries
 43      fact_patterns = [r'what is', r'who is', r'when was', r'where is', r'capital of']
 44      for pattern in fact_patterns:
 45          if re.search(pattern, query_lower):
 46              return f"Answer this question directly and concisely: {query}"
 47      
 48      # Technical/explanation queries
 49      if any(word in query_lower for word in ['explain', 'how does', 'describe', 'what are']):
 50          return f"Explain in detail: {query}. Be technical and comprehensive."
 51      
 52      return query
 53  
 54  # ====================== API HANDLERS ======================
 55  async def try_openrouter(session: aiohttp.ClientSession, query: str) -> Dict[str, Any]:
 56      """Endpoint 1: OpenRouter"""
 57      if not OPENROUTER_API_KEY:
 58          return {"error": "No API key", "success": False}
 59      
 60      headers = {
 61          "Authorization": f"Bearer {OPENROUTER_API_KEY}",
 62          "Content-Type": "application/json"
 63      }
 64      
 65      payload = {
 66          "model": OPENROUTER_MODEL,
 67          "messages": [{"role": "user", "content": query}],
 68          "temperature": 0.7,
 69          "max_tokens": 1000
 70      }
 71      
 72      try:
 73          async with session.post(OPENROUTER_URL, json=payload, headers=headers) as resp:
 74              if resp.status == 200:
 75                  data = await resp.json()
 76                  return {
 77                      "success": True,
 78                      "response": data['choices'][0]['message']['content'],
 79                      "source": "openrouter"
 80                  }
 81              else:
 82                  return {"error": f"HTTP {resp.status}", "success": False}
 83      except Exception as e:
 84          return {"error": str(e), "success": False}
 85  
 86  async def try_groq(session: aiohttp.ClientSession, query: str) -> Dict[str, Any]:
 87      """Endpoint 2: Groq"""
 88      if not GROQ_API_KEY:
 89          return {"error": "No API key", "success": False}
 90      
 91      headers = {
 92          "Authorization": f"Bearer {GROQ_API_KEY}",
 93          "Content-Type": "application/json"
 94      }
 95      
 96      payload = {
 97          "model": GROQ_MODEL,
 98          "messages": [{"role": "user", "content": query}],
 99          "temperature": 0.7,
100          "max_tokens": 1000
101      }
102      
103      try:
104          async with session.post(GROQ_URL, json=payload, headers=headers) as resp:
105              if resp.status == 200:
106                  data = await resp.json()
107                  return {
108                      "success": True,
109                      "response": data['choices'][0]['message']['content'],
110                      "source": "groq"
111                  }
112              else:
113                  return {"error": f"HTTP {resp.status}", "success": False}
114      except Exception as e:
115          return {"error": str(e), "success": False}
116  
117  async def try_ollama(session: aiohttp.ClientSession, query: str) -> Dict[str, Any]:
118      """Endpoint 3: Local Ollama"""
119      payload = {
120          "model": OLLAMA_MODEL,
121          "prompt": query,
122          "stream": False,
123          "options": {"temperature": 0.7}
124      }
125      
126      try:
127          async with session.post(OLLAMA_URL, json=payload, timeout=30) as resp:
128              if resp.status == 200:
129                  data = await resp.json()
130                  return {
131                      "success": True,
132                      "response": data.get('response', ''),
133                      "source": "ollama"
134                  }
135              else:
136                  return {"error": f"HTTP {resp.status}", "success": False}
137      except Exception as e:
138          return {"error": str(e), "success": False}
139  
140  async def try_custom_api(session: aiohttp.ClientSession, query: str) -> Dict[str, Any]:
141      """Endpoint 4: Custom/io.net style API"""
142      headers = {"Content-Type": "application/json"}
143      if CUSTOM_API_KEY:
144          headers["Authorization"] = f"Bearer {CUSTOM_API_KEY}"
145      
146      payload = {
147          "model": CUSTOM_MODEL,
148          "messages": [{"role": "user", "content": query}],
149          "temperature": 0.7,
150          "max_tokens": 1000
151      }
152      
153      try:
154          async with session.post(CUSTOM_API_URL, json=payload, headers=headers, timeout=60) as resp:
155              if resp.status == 200:
156                  data = await resp.json()
157                  return {
158                      "success": True,
159                      "response": data['choices'][0]['message']['content'],
160                      "source": "custom"
161                  }
162              else:
163                  return {"error": f"HTTP {resp.status}", "success": False}
164      except Exception as e:
165          return {"error": str(e), "success": False}
166  
167  # ====================== ORCHESTRATOR ======================
168  async def get_ai_response(query: str, mode: str = "auto") -> Dict[str, Any]:
169      """Main orchestrator with 4 endpoints"""
170      enhanced_query = enhance_prompt(query)
171      print(f"Enhanced query: {enhanced_query[:50]}...")
172      
173      async with aiohttp.ClientSession() as session:
174          start_time = datetime.now()
175          
176          if mode == "cloud":
177              # Try cloud endpoints only
178              print("  Trying OpenRouter...")
179              result = await try_openrouter(session, enhanced_query)
180              if result['success']:
181                  result['time'] = (datetime.now() - start_time).total_seconds()
182                  return result
183              
184              print("  OpenRouter failed, trying Groq...")
185              result = await try_groq(session, enhanced_query)
186              result['time'] = (datetime.now() - start_time).total_seconds()
187              return result
188          
189          elif mode == "local":
190              # Try local endpoints only
191              print("  Trying Ollama...")
192              result = await try_ollama(session, enhanced_query)
193              if result['success']:
194                  result['time'] = (datetime.now() - start_time).total_seconds()
195                  return result
196              
197              print("  Ollama failed, trying Custom API...")
198              result = await try_custom_api(session, enhanced_query)
199              result['time'] = (datetime.now() - start_time).total_seconds()
200              return result
201          
202          else:  # auto mode
203              # Try all endpoints in order
204              endpoints = [
205                  ("OpenRouter", try_openrouter),
206                  ("Groq", try_groq),
207                  ("Ollama", try_ollama),
208                  ("Custom API", try_custom_api)
209              ]
210              
211              for name, endpoint_func in endpoints:
212                  print(f"  Trying {name}...")
213                  result = await endpoint_func(session, enhanced_query)
214                  if result['success']:
215                      result['time'] = (datetime.now() - start_time).total_seconds()
216                      return result
217              
218              # All failed
219              result['time'] = (datetime.now() - start_time).total_seconds()
220              return result
221  
222  # ====================== MAIN ======================
223  async def main():
224      if len(sys.argv) < 2:
225          print("Usage: python3 ai_orchestrator_enhanced.py 'your query' [--mode auto|cloud|local]")
226          print("Modes: auto (default), cloud (OpenRouter/Groq), local (Ollama/Custom)")
227          sys.exit(1)
228      
229      # Parse arguments
230      args = sys.argv[1:]
231      mode = "auto"
232      query_parts = []
233      
234      for arg in args:
235          if arg.startswith("--mode="):
236              mode = arg.split("=")[1]
237          elif arg == "--cloud":
238              mode = "cloud"
239          elif arg == "--local":
240              mode = "local"
241          elif arg.startswith("--"):
242              continue
243          else:
244              query_parts.append(arg)
245      
246      query = " ".join(query_parts)
247      
248      print("=" * 60)
249      print(f"Query: {query}")
250      print(f"Mode: {mode}")
251      print("=" * 60)
252      
253      result = await get_ai_response(query, mode)
254      
255      if result['success']:
256          print("✅ SUCCESS")
257          print("=" * 60)
258          print(f"Source: {result['source']}")
259          print(f"Time: {result['time']:.1f}s")
260          print("\n🎯 ANSWER:")
261          print("=" * 60)
262          print(result['response'])
263          print("=" * 60)
264          
265          # Save result
266          timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
267          filename = f"ai_result_{timestamp}.json"
268          with open(filename, "w") as f:
269              json.dump({
270                  "query": query,
271                  "response": result['response'],
272                  "source": result['source'],
273                  "time": result['time'],
274                  "mode": mode,
275                  "success": True
276              }, f, indent=2)
277          print(f"\n💾 Saved to: {filename}")
278      else:
279          print("❌ ALL ENDPOINTS FAILED")
280          print("=" * 60)
281          print(f"Error: {result.get('error', 'Unknown error')}")
282          print(f"Time: {result.get('time', 0):.1f}s")
283  
284  if __name__ == "__main__":
285      asyncio.run(main())