support for OpenRouter and LiteLLM as options, adding model provider option

Author: joshpxyne

gpt_migrate/ai.py

@@ -1,4 +1,3 @@
-from langchain.chat_models import ChatOpenAI
 import os
 import openai
 from utils import parse_code_string
@@ -8,25 +7,35 @@
 openai.api_key = os.getenv("OPENROUTER_API_KEY")
 
 class AI:
-    def __init__(self, model="gpt-4-32k", temperature=0.1, max_tokens=10000):
+    def __init__(self, model="gpt-4-32k", model_provider="openai", modelrouter="openrouter", temperature=0.1, max_tokens=10000):
         self.temperature = temperature
         self.max_tokens = max_tokens
+        self.model_provider = model_provider
         self.model_name = model
-        try:
-            _ = ChatOpenAI(model_name=model) # check to see if model is available to user
-        except Exception as e:
-            print(e)
-            self.model_name = "gpt-3.5-turbo"
+        self.modelrouter = modelrouter
 
     def write_code(self, prompt):
         message=[{"role": "user", "content": str(prompt)}] 
-        response = completion(
-            messages=message,
-            stream=False,
-            model=self.model_name,
-            max_tokens=self.max_tokens,
-            temperature=self.temperature
-        )
+        if self.modelrouter == "openrouter":
+            response = openai.ChatCompletion.create(
+                model="{}/{}".format(self.model_provider,self.model_name), # Optional (user controls the default)
+                messages=message,
+                stream=False,
+                max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                headers={
+                    "HTTP-Referer": "https://gpt-migrate.com",
+                    "X-Title": "GPT-Migrate",
+                },
+            )
+        else:
+            response = completion(
+                messages=message,
+                stream=False,
+                model=self.model_name,
+                max_tokens=self.max_tokens,
+                temperature=self.temperature
+            )
         if response["choices"][0]["message"]["content"].startswith("INSTRUCTIONS:"):
             return ("INSTRUCTIONS:","",response["choices"][0]["message"]["content"][14:])
         else:
@@ -35,17 +44,32 @@ def write_code(self, prompt):
 
     def run(self, prompt):
         message=[{"role": "user", "content": str(prompt)}] 
-        response = completion(
-            messages=message,
-            stream=True,
-            model=self.model_name,
-            max_tokens=self.max_tokens,
-            temperature=self.temperature
-        )
-        chat = ""
-        for chunk in response:
-            delta = chunk["choices"][0]["delta"]
-            msg = delta.get("content", "")
-            chat += msg
-        return chat
+        if self.modelrouter == "openrouter":
+            response = openai.ChatCompletion.create(
+                model="{}/{}".format(self.model_provider,self.model_name), # Optional (user controls the default)
+                messages=message,
+                stream=False,
+                max_tokens=self.max_tokens,
+                temperature=self.temperature,
+                headers={
+                    "HTTP-Referer": "https://gpt-migrate.com",
+                    "X-Title": "GPT-Migrate",
+                },
+            )
+            return response["choices"][0]["message"]["content"]
+        else:
+            response = completion(
+                messages=message,
+                stream=True,
+                model=self.model_name,
+                max_tokens=self.max_tokens,
+                temperature=self.temperature
+            )
+            chat = ""
+            for chunk in response:
+                delta = chunk["choices"][0]["delta"]
+                msg = delta.get("content", "")
+                chat += msg
+            return chat
+        
 

gpt_migrate/main.py

@@ -32,6 +32,8 @@ def __init__(self, sourcedir, targetdir, sourcelang, targetlang, sourceentry, so
 @app.command()
 def main(
         model: str = typer.Option("gpt-4-32k", help="Large Language Model to be used."),
+        model_provider: str = typer.Option("openai", help="Model provider to be used."),
+        modelrouter: str = typer.Option("openrouter", help="Model router to be used. Options are 'openrouter' or 'litellm'."),
         temperature: float = typer.Option(0, help="Temperature setting for the AI model."),
         sourcedir: str = typer.Option("../benchmarks/flask-nodejs/source", help="Source directory containing the code to be migrated."),
         sourcelang: str = typer.Option(None, help="Source language or framework of the code to be migrated."),
@@ -48,7 +50,9 @@ def main(
 
     ai = AI(
         model=model,
+        model_provider=model_provider,
         temperature=temperature,
+        modelrouter=modelrouter
     )
 
     sourcedir = os.path.abspath(sourcedir)

gpt_migrate/memory/external_dependencies

@@ -1,6 +1,3 @@
-rocket
-serde
-serde_json
+express
 bcrypt
-rusqlite
-serde_json
+node-json-db

gpt_migrate/memory/gpt_migrate/db.js_sigs.json

@@ -0,0 +1 @@
+[{"signature": "readItems() -> object", "description": "Reads the items from a JSON file and returns them as an object."}, {"signature": "writeItems(groceryItems: object)", "description": "Writes the given grocery items object to a JSON file."}]