wip

app.py

@@ -1,9 +1,10 @@
 from flask import Flask, request, jsonify, render_template, make_response, send_from_directory
-import anthropic
+from litellm import completion
 import os
 import json
 from datetime import datetime
 import base64
+from llm_utils import generate_completion
 
 app = Flask(__name__)
 
@@ -66,27 +67,16 @@ def open_pdf(filename):
 def generate_flashcard():
     data = request.json
     prompt = data['prompt']
-    api_key = request.headers.get('X-API-Key')
     mode = data.get('mode', 'flashcard')
 
-    client = anthropic.Anthropic(api_key=api_key)
-
     try:
-        model = data.get('model', "claude-3-5-sonnet-20240620")
-        message = client.messages.create(
-            model=model,
-            max_tokens=1024,
-            messages=[
-                {"role": "user", "content": prompt}
-            ]
-        )
-        
-        content = message.content[0].text
+        # Use llm_utils to generate completion
+        content = generate_completion(prompt)
         print(prompt)
         print(content)
 
         if mode == 'language':
-            # For Language mode, parse the content and return in the custom format
+            # Parse language learning format
             lines = content.split('\n')
             word = ''
             translation = ''
@@ -106,8 +96,10 @@ def generate_flashcard():
                 'translation': translation,
                 'answer': answer
             }
-            response = make_response(jsonify({'flashcard': flashcard}))
+            return jsonify({'flashcard': flashcard})
+            
         elif mode == 'flashcard' or 'flashcard' in prompt.lower():
+            # Parse flashcard format
             flashcards = []
             current_question = ''
             current_answer = ''
@@ -124,17 +116,15 @@ def generate_flashcard():
             if current_question and current_answer:
                 flashcards.append({'question': current_question, 'answer': current_answer})
 
-            response = make_response(jsonify({'flashcards': flashcards}))
+            return jsonify({'flashcards': flashcards})
+            
         elif mode == 'explain' or 'explain' in prompt.lower():
-            # For Explain mode, return the entire content as the explanation
-            response = make_response(jsonify({'explanation': content}))
+            # Return explanation format
+            return jsonify({'explanation': content})
+            
         else:
-            response = make_response(jsonify({'error': 'Invalid mode'}))
-
-        # Set cookie with the API key
-        response.set_cookie('last_working_api_key', api_key, secure=True, httponly=True, samesite='Strict')
-
-        return response
+            return jsonify({'error': 'Invalid mode'}), 400
+        
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 if __name__ == '__main__':

llm_utils.py

@@ -0,0 +1,40 @@
+from litellm import completion
+import yaml
+import os
+
+def load_model_config():
+    with open('models.yaml', 'r') as file:
+        return yaml.safe_load(file)
+
+def generate_completion(prompt: str, api_key: str = None) -> str:
+    """
+    Generate completion using LiteLLM with the configured model
+    
+    Args:
+        prompt (str): The input prompt
+        api_key (str, optional): Override API key. If not provided, will use environment variable
+        
+    Returns:
+        str: The generated completion text
+    """
+    config = load_model_config()
+    
+    # Get the first environment variable and its models
+    first_env_var = list(config['models'][0].keys())[0]
+    model_name = config['models'][0][first_env_var][0]
+    
+    # If no API key provided, get from environment
+    if api_key is None:
+        api_key = os.getenv(first_env_var)
+        if not api_key:
+            raise ValueError(f"Please set {first_env_var} environment variable")
+    
+    messages = [{"role": "user", "content": prompt}]
+    
+    response = completion(
+        model=model_name,
+        messages=messages,
+        api_key=api_key
+    )
+    
+    return response.choices[0].message.content 

models.yaml

@@ -0,0 +1,10 @@
+models:
+  - GEMINI_API_KEY:
+      - "gemini/gemini-exp-1206"
+
+  - OPENROUTER_API_KEY:
+      - "openrouter/google/gemini-exp-1206:free"
+      - "openrouter/anthropic/claude-3-haiku-20240307"
+      - "openrouter/anthropic/claude-3-sonnet-20240229"
+  
+  

templates/index.html

@@ -9,6 +9,45 @@
     <script src="https://cdnjs.cloudflare.com/ajax/libs/jszip/3.1.5/jszip.min.js"></script>
     <script src="https://cdn.jsdelivr.net/npm/epubjs/dist/epub.min.js"></script>
     <link rel="stylesheet" href="/static/css/styles.css">
+    <style>
+        .api-settings {
+            margin-bottom: 15px;
+            padding: 10px;
+            background: #f5f5f5;
+            border-radius: 5px;
+        }
+
+        .model-group {
+            margin-bottom: 10px;
+        }
+
+        #model-select {
+            width: 100%;
+            padding: 8px;
+            margin-bottom: 10px;
+            border: 1px solid #ddd;
+            border-radius: 4px;
+        }
+
+        #custom-model-inputs {
+            margin-top: 10px;
+        }
+
+        #custom-model-inputs input {
+            width: 100%;
+            padding: 8px;
+            margin-bottom: 5px;
+            border: 1px solid #ddd;
+            border-radius: 4px;
+        }
+
+        .default-api-key input {
+            width: 100%;
+            padding: 8px;
+            border: 1px solid #ddd;
+            border-radius: 4px;
+        }
+    </style>
 </head>
 
 <body>

test_llm_utils.py

@@ -0,0 +1,40 @@
+import os
+import yaml
+from llm_utils import generate_completion
+
+def test_generate_completion():
+    # Load config to get the environment variable name
+    with open('models.yaml', 'r') as file:
+        config = yaml.safe_load(file)
+    
+    # Get the first environment variable name
+    env_var_name = list(config['models'][0].keys())[0]
+    
+    # Get API key from environment variable
+    api_key = os.getenv(env_var_name)
+    if not api_key:
+        raise ValueError(f"Please set {env_var_name} environment variable")
+    
+    # Test prompt
+    test_prompt = "What is 2+2? Answer in one word."
+    
+    try:
+        # Test with explicit API key
+        response = generate_completion(test_prompt, api_key)
+        print(f"Test prompt: {test_prompt}")
+        print(f"Response with explicit API key: {response}")
+        assert isinstance(response, str)
+        assert len(response) > 0
+        
+        # Test with environment variable
+        response = generate_completion(test_prompt)
+        print(f"Response with environment variable: {response}")
+        assert isinstance(response, str)
+        assert len(response) > 0
+        
+        print("Test passed successfully!")
+    except Exception as e:
+        print(f"Test failed with error: {str(e)}")
+
+if __name__ == "__main__":
+    test_generate_completion()