query.py

#!/usr/bin/env python3
import argparse
import configparser
import sqlite3
import openai
import pinecone
config = configparser.ConfigParser()
config.read("config.ini")


def openai_embeddings(content: str, model: str = "text-embedding-ada-002"):
    res = openai.Embedding.create(model=model, input=content)
    vec = res["data"][0]["embedding"]  # type: ignore
    return vec

def main(query):
    vec = openai_embeddings(query)


    pinecone.init(api_key=config["PINECONE"]["api"], environment="gcp-starter")
    index = pinecone.Index("tech-articles")
    resp = index.query(vec, top_k=3)

    matches = [item['id'] for item in resp['matches']]

    DB: str = config["SOURCE"]["db"]
    TABLE: str = config["SOURCE"]["table"]
    conn = sqlite3.connect(DB)
    cursor = conn.cursor()

    placeholders = ', '.join(['?'] * len(matches))
    sqlite_q = f"SELECT content FROM tech_articles WHERE sign IN ({placeholders})"
    cursor.execute(sqlite_q, matches)
    result = cursor.fetchall()

    for row in result:
        print(row)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fetch summaries based on query.")
    parser.add_argument("query", type=str, help="Query to search for.")
    args = parser.parse_args()
    
    main(args.query)