summaryrefslogtreecommitdiffstats
path: root/g4f/gui/server/backend.py
blob: e24d4da2e61c2425c818846fbfbce1cd388efba2 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import json
import os
from flask import request, Flask, jsonify, send_from_directory
from g4f.image import is_allowed_extension, to_image
from .api import Api

class Backend_Api(Api):    
    """
    Handles various endpoints in a Flask application for backend operations.

    This class provides methods to interact with models, providers, and to handle
    various functionalities like conversations, error handling, and version management.

    Attributes:
        app (Flask): A Flask application instance.
        routes (dict): A dictionary mapping API endpoints to their respective handlers.
    """
    def __init__(self, app: Flask) -> None:
        """
        Initialize the backend API with the given Flask application.

        Args:
            app (Flask): Flask application instance to attach routes to.
        """
        self.app: Flask = app
        self.routes = {
            '/backend-api/v2/models': {
                'function': self.get_models,
                'methods': ['GET']
            },
            '/backend-api/v2/models/<provider>': {
                'function': self.get_provider_models,
                'methods': ['GET']
            },
            '/backend-api/v2/image_models': {
                'function': self.get_image_models,
                'methods': ['GET']
            },
            '/backend-api/v2/providers': {
                'function': self.get_providers,
                'methods': ['GET']
            },
            '/backend-api/v2/version': {
                'function': self.get_version,
                'methods': ['GET']
            },
            '/backend-api/v2/conversation': {
                'function': self.handle_conversation,
                'methods': ['POST']
            },
            '/backend-api/v2/error': {
                'function': self.handle_error,
                'methods': ['POST']
            },
            '/images/<path:name>': {
                'function': self.serve_images,
                'methods': ['GET']
            },
            '/images': {
                'function': self.get_images,
                'methods': ['GET']
            }
        }

    def handle_error(self):
        """
        Initialize the backend API with the given Flask application.

        Args:
            app (Flask): Flask application instance to attach routes to.
        """
        print(request.json)
        return 'ok', 200

    def handle_conversation(self):
        """
        Handles conversation requests and streams responses back.

        Returns:
            Response: A Flask response object for streaming.
        """
        
        kwargs = {}
        if "file" in request.files:
            file = request.files['file']
            if file.filename != '' and is_allowed_extension(file.filename):
                kwargs['image'] = to_image(file.stream, file.filename.endswith('.svg'))
                kwargs['image_name'] = file.filename
        if "json" in request.form:
            json_data = json.loads(request.form['json'])
        else:
            json_data = request.json

        kwargs = self._prepare_conversation_kwargs(json_data, kwargs)

        return self.app.response_class(
            self._create_response_stream(kwargs, json_data.get("conversation_id"), json_data.get("provider")),
            mimetype='text/event-stream'
        )

    def get_provider_models(self, provider: str):
        models = super().get_provider_models(provider)
        if models is None:
            return 404, "Provider not found"
        return models

    def _format_json(self, response_type: str, content) -> str:
        """
        Formats and returns a JSON response.

        Args:
            response_type (str): The type of the response.
            content: The content to be included in the response.

        Returns:
            str: A JSON formatted string.
        """
        return json.dumps(super()._format_json(response_type, content)) + "\n"

    @staticmethod
    def get_images():
        images_dir = "./generated_images"
        try:
            images = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]
            images = [f"/images/{image}" for image in images if image.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp'))]
            return jsonify(images)
        except Exception as e:
            return str(e), 500

    @staticmethod
    def serve_images(name):
        images_dir = "./generated_images"
        return send_from_directory(os.path.abspath(images_dir), name)