File: cli.py

package info (click to toggle)
redis 5%3A8.0.2-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 22,304 kB
  • sloc: ansic: 216,903; tcl: 51,562; sh: 4,625; perl: 4,214; cpp: 3,568; python: 2,954; makefile: 2,055; ruby: 639; javascript: 30; csh: 7
file content (147 lines) | stat: -rwxr-xr-x 4,914 bytes parent folder | download
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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
#
# Copyright (c) 2009-Present, Redis Ltd.
# All rights reserved.
#
# Licensed under your choice of (a) the Redis Source Available License 2.0
# (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the
# GNU Affero General Public License v3 (AGPLv3).
#

#!/usr/bin/env python3
import redis
import requests
import re
import shlex
from prompt_toolkit import PromptSession
from prompt_toolkit.history import InMemoryHistory

def get_embedding(text):
    """Get embedding from local Ollama API"""
    url = "http://localhost:11434/api/embeddings"
    payload = {
        "model": "mxbai-embed-large",
        "prompt": text
    }
    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json()['embedding']
    except requests.exceptions.RequestException as e:
        raise Exception(f"Failed to get embedding: {str(e)}")

def process_embedding_patterns(text):
    """Process !"text" and !!"text" patterns in the command"""

    def replace_with_embedding(match):
        text = match.group(1)
        embedding = get_embedding(text)
        return f"VALUES {len(embedding)} {' '.join(map(str, embedding))}"

    def replace_with_embedding_and_text(match):
        text = match.group(1)
        embedding = get_embedding(text)
        # Return both the embedding values and the original text as next argument
        return f'VALUES {len(embedding)} {" ".join(map(str, embedding))} "{text}"'

    # First handle !!"text" pattern (must be done before !"text")
    text = re.sub(r'!!"([^"]*)"', replace_with_embedding_and_text, text)
    # Then handle !"text" pattern
    text = re.sub(r'!"([^"]*)"', replace_with_embedding, text)
    return text

def parse_command(command):
    """Parse command respecting quoted strings"""
    try:
        # Use shlex to properly handle quoted strings
        return shlex.split(command)
    except ValueError as e:
        raise Exception(f"Invalid command syntax: {str(e)}")

def format_response(response):
    """Format the response to match Redis protocol style"""
    if response is None:
        return "(nil)"
    elif isinstance(response, bool):
        return "+OK" if response else "(error) Operation failed"
    elif isinstance(response, (list, set)):
        if not response:
            return "(empty list or set)"
        return "\n".join(f"{i+1}) {item}" for i, item in enumerate(response))
    elif isinstance(response, int):
        return f"(integer) {response}"
    else:
        return str(response)

def main():
    # Default connection to localhost:6379
    r = redis.Redis(host='localhost', port=6379, decode_responses=True)

    try:
        # Test connection
        r.ping()
        print("Connected to Redis. Type your commands (CTRL+D to exit):")
        print("Special syntax:")
        print("  !\"text\"  - Replace with embedding")
        print("  !!\"text\" - Replace with embedding and append text as value")
        print("  \"text\"   - Quote strings containing spaces")
    except redis.ConnectionError:
        print("Error: Could not connect to Redis server")
        return

    # Setup prompt session with history
    session = PromptSession(history=InMemoryHistory())

    # Main loop
    while True:
        try:
            # Read input with line editing support
            command = session.prompt("redis> ")

            # Skip empty commands
            if not command.strip():
                continue

            # Process any embedding patterns before parsing
            try:
                processed_command = process_embedding_patterns(command)
            except Exception as e:
                print(f"(error) Embedding processing failed: {str(e)}")
                continue

            # Parse the command respecting quoted strings
            try:
                parts = parse_command(processed_command)
            except Exception as e:
                print(f"(error) {str(e)}")
                continue

            if not parts:
                continue

            cmd = parts[0].lower()
            args = parts[1:]

            # Execute command
            try:
                method = getattr(r, cmd, None)
                if method is not None:
                    result = method(*args)
                else:
                    # Use execute_command for unknown commands
                    result = r.execute_command(cmd, *args)
                print(format_response(result))
            except AttributeError:
                print(f"(error) Unknown command '{cmd}'")

        except EOFError:
            print("\nGoodbye!")
            break
        except KeyboardInterrupt:
            continue  # Allow Ctrl+C to clear current line
        except redis.RedisError as e:
            print(f"(error) {str(e)}")
        except Exception as e:
            print(f"(error) {str(e)}")

if __name__ == "__main__":
    main()