omi/backend/scripts/rag/_shared.py at main · ghzhost/omi · GitHub

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
import json
import os

# noinspection PyUnresolvedReferences
from typing import Any, List, cast

# noinspection PyUnresolvedReferences
import plotly.graph_objects as go

from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings
from pinecone import Pinecone

load_dotenv('../../.env')
os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '../../' + os.getenv('GOOGLE_APPLICATION_CREDENTIALS', '')

index: Any
if os.getenv('PINECONE_API_KEY') is not None:
    pc = cast(Any, Pinecone(api_key=os.getenv('PINECONE_API_KEY', '')))
    index = pc.Index(os.getenv('PINECONE_INDEX_NAME', ''))
else:
    index = None

import database.conversations as conversations_db

uid = 'viUv7GtdoHXbK1UBCDlPuTDuPgJ2'

openai_embeddings = OpenAIEmbeddings(model="text-embedding-3-large")


def query_vectors(query: str, uid: str, k: int = 1000) -> List[List[Any]]:
    xq = openai_embeddings.embed_query(query)
    xc = index.query(vector=xq, top_k=k, filter={'uid': uid}, namespace="ns1", include_values=True)
    data: List[List[Any]] = []
    for item in xc['matches']:
        data.append([item['id'].replace(f'{uid}-', ''), item['values']])
    print('Found:', len(data), 'vectors')
    return data


def get_memories(ignore_cached: bool = False):
    if not os.path.exists('memories.json') or ignore_cached:
        memories = conversations_db.get_conversations(uid, limit=1000)
        if ignore_cached:
            return memories

        with open('memories.json', 'w') as f:
            f.write(json.dumps(memories, indent=4, default=str))

    with open('memories.json', 'r') as f:
        return json.loads(f.read())


def get_all_markers(data: List[List[Any]], data_points: Any, target: int) -> Any:
    return go.Scatter(
        x=data_points[target:, 0],
        y=data_points[target:, 1],
        mode='markers',
        marker=dict(size=8, opacity=0.5, color='blue'),
        text=[f"{item[0]}" for item in data[5:]],
        hoverinfo='text',
        name='Other Memories',
    )


def get_top_markers(data: List[List[Any]], data_points: Any, target: int) -> Any:
    return go.Scatter(
        x=data_points[:target, 0],
        y=data_points[:target, 1],
        mode='markers',
        marker=dict(size=10, opacity=0.8, color='green'),
        text=[f"Top {i + 1}: {item[0]}" for i, item in enumerate(data[:5])],
        hoverinfo='text',
        name='Top Matches',
    )


def get_query_marker(query_point: Any, query: str) -> Any:
    return go.Scatter(
        x=[query_point[0]],
        y=[query_point[1]],
        mode='markers',
        marker=dict(symbol='x', size=12, color='red', line=dict(width=2)),
        text=[query],
        hoverinfo='text',
        name='Query',
    )


def generate_html_visualization(fig: Any, file_name: str = 'embedding_visualization.html') -> None:
    fig.update_layout(
        title=f'Embedding Visualization',
        xaxis_title='UMAP Dimension 1',
        yaxis_title='UMAP Dimension 2',
        width=800,
        height=600,
        showlegend=True,
    )

    # Generate HTML
    html_content = f'''
        <html>
            <head>
                <title>Embedding Visualization</title>
                <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
            </head>
            <body>
                <div id="plotDiv"></div>
                <script>
                    var plotlyData = {fig.to_json()};
                    Plotly.newPlot('plotDiv', plotlyData.data, plotlyData.layout);
                </script>
            </body>
        </html>
        '''

    with open(file_name, 'w') as f:
        f.write(html_content)

    print(f"HTML file '{file_name}' has been generated.")