57 lines
1.5 KiB
Python
Executable File
57 lines
1.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
|
|
from index_common import embed_text, get_conn, vector_literal
|
|
|
|
|
|
def sanitize(text: str) -> str:
|
|
return text.replace('```', '` ` `').strip()
|
|
|
|
|
|
def query(text: str, top_k: int = 5, max_chars: int = 2500) -> str:
|
|
embedding = embed_text(text)
|
|
vector = vector_literal(embedding)
|
|
|
|
with get_conn() as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'''
|
|
SELECT source, content
|
|
FROM memory_primary
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT %s
|
|
''',
|
|
(vector, top_k),
|
|
)
|
|
rows = cur.fetchall()
|
|
|
|
parts: list[str] = []
|
|
total_len = 0
|
|
for source, content in rows:
|
|
snippet = sanitize(content[:600])
|
|
block = f'<retrieved_context source="{source}">\n{snippet}\n</retrieved_context>'
|
|
if total_len + len(block) > max_chars:
|
|
break
|
|
parts.append(block)
|
|
total_len += len(block)
|
|
|
|
return '\n\n'.join(parts)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('query', nargs='+')
|
|
parser.add_argument('--top-k', type=int, default=5)
|
|
parser.add_argument('--max-chars', type=int, default=2500)
|
|
args = parser.parse_args()
|
|
|
|
text = ' '.join(args.query).strip()
|
|
if not text:
|
|
return
|
|
print(query(text, top_k=args.top_k, max_chars=args.max_chars))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|