¡Hola developer 👋🏻! En el último vídeo de mi serie sobre IA Generativa te mostré una API en Flask con la que era capaz de llamar tanto a modelos que se estaban ejecutando gracias a Ollama o alojados en GitHub Models.

Si bien es cierto que durante el vídeo, con efectos educativos, no refactoricé del todo la misma, lo prometido es deuda, y aquí te dejo una versión más limpia de la API para que veas lo sencillo que es llamar a diferentes modelos a través de una sola API y usando el SDK de OpenAI para todo.

requirements.txt

Para que la siguiente implementación que te comparto funcione como se espera, es necesario tener un archivo requirements.txt como el siguiente:

flask==3.1.0
flask-cors==5.0.1
python-dotenv==1.1.0
mistralai==1.6.0
openai==1.70.0
requests==2.31.0
tiktoken==0.9.0

La API

from flask import Flask, Response, request, jsonify
from flask_cors import CORS
from openai import OpenAI
import os
from dotenv import load_dotenv
import tiktoken
import prompty.openai
# Load environment variables from a .env file
load_dotenv()
app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

def create_openai_client(source):
    """Create an OpenAI client based on the source"""
    if source == 'github':
        return OpenAI(
            base_url=os.getenv("GITHUB_MODELS_URL"),
            api_key=os.getenv("GITHUB_TOKEN"),
        )
    elif source == 'ollama':
        ollama_url = os.getenv("OLLAMA_URL")
        return OpenAI(
            base_url=f"{ollama_url}/v1",
            api_key="ollama",  # Ollama doesn't require a real API key, but one is needed for the SDK
        )
    return None

def process_stream_response(stream_response):
    """Process the streaming response from the API"""
    for chunk in stream_response:
        if chunk.choices and len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None:
            yield chunk.choices[0].delta.content

@app.route("/generate")
def generate():
    model_name = request.args.get('model')
    title = request.args.get('title')    
    source = request.args.get('source')
    print(f"Source: {source}")
    print(f"Model: {model_name}")
    print(f"Title: {title}")
    def generate_stream():
        if source not in ['github', 'ollama']:
            yield f"Unknown source: {source}\n"
            return
        try:
            client = create_openai_client(source)
            if not client:
                yield f"Failed to create client for source: {source}"
                return
                
            stream_response = client.chat.completions.create(
                messages=[
                    {"role": "user", "content": f"Mejorame el siguiente titulo, incluye emojis: '{title}'"}
                ],
                model=model_name,
                stream=True
            )
            
            yield from process_stream_response(stream_response)
                
        except Exception as e:
            yield f"Error using OpenAI SDK with {source}: {str(e)}"
    return Response(generate_stream(), content_type="text/event-stream")

@app.route("/count_tokens", methods=["POST"])
def count_tokens():
    data = request.json
    if not data or 'text' not in data:
        return jsonify({"error": "No text provided"}), 400
    text = data['text']
    try:
        # Get the CL100K base encoding
        encoding = tiktoken.get_encoding("cl100k_base")
        tokens = encoding.encode(text)
        token_count = len(tokens)
        # Create a list to hold the token representations
        token_representations = []
        for token in tokens:
            # Decode each token individually to its text representation
            token_text = encoding.decode([token])
            token_representations.append({
                "token_id": int(token),
                "token_text": token_text
            })
        print(f"Token count: {token_count}")
        return jsonify({
            "token_count": token_count,
            "tokens": token_representations
        }), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

Por último, necesitas un .env con los siguientes valores:

GITHUB_TOKEN="github_pat_XXXX"
GITHUB_MODELS_URL="https://models.github.ai/inference"
OLLAMA_URL="http://host.docker.internal:11434"

Si además no quieres tener que instalar absolutamente nada en tu máquina local, puedes crear una configuración de dev containers como esta:

{
	"name": "Hoy empiezo con IA Generativa 📖🤖",
	"image": "mcr.microsoft.com/devcontainers/python:1-3.12-bullseye",
	"features": {
		"ghcr.io/prulloac/devcontainer-features/ollama:1": {}
	},
	"customizations": {
		"vscode": {
			"extensions": [
				"ms-windows-ai-studio.windows-ai-studio",
				"GitHub.copilot",
				"GitHub.copilot-chat",
				"ms-toolsai.prompty",
				"ritwickdey.LiveServer"
			]
		}
	},
	"hostRequirements": {
		"memory": "16gb",
		"cpus": 4
	},
	{
		"postCreateCommand": "pip install --upgrade pip && pip install -r requirements.txt"
	}
}

y si quieres saber más sobre Dev Containers puedes echarle un vistazo a este otro vídeo 😇

¡Nos vemos 👋🏻!