From 98846d8165057b85583bcd9fc7fecf5a46010664 Mon Sep 17 00:00:00 2001
From: Yura Lukashik <YuraLukashik@gmail.com>
Date: Sun, 29 Oct 2023 19:03:46 +0300
Subject: [PATCH] Search: try custom HTTPS implementation for Pinecone request

---
 .../external_services/pinecone.py             | 20 +++++++
 .../semantic_search/https_requests.py         | 53 +++++++++++++++++++
 src/semantic_search/semantic_search/query.py  | 18 +++++--
 3 files changed, 86 insertions(+), 5 deletions(-)
 create mode 100644 src/semantic_search/semantic_search/https_requests.py

diff --git a/src/semantic_search/semantic_search/external_services/pinecone.py b/src/semantic_search/semantic_search/external_services/pinecone.py
index 1989b23..a399909 100644
--- a/src/semantic_search/semantic_search/external_services/pinecone.py
+++ b/src/semantic_search/semantic_search/external_services/pinecone.py
@@ -1,6 +1,8 @@
 import logging
 
 import pinecone
+
+from ..https_requests import send_https_request
 from ..config import get_pinecone_key, get_pinecone_environment, get_pinecone_index_name
 
 pinecone.init(api_key=get_pinecone_key(), environment=get_pinecone_environment(), log_level="debug")
@@ -12,3 +14,21 @@
 
 def get_pinecone_index() -> 'pinecone.Index':
     return pinecone.Index(get_pinecone_index_name())
+
+
+def query_index(query_vector, top_k, namespace, include_values, include_metadata):
+    host = f"{get_pinecone_index_name()}-0ddc4d6.svc.{get_pinecone_environment()}.pinecone.io"
+    path = "/query"
+    headers = {
+        "content-type": "application/json",
+        "api-key": get_pinecone_key(),
+        "accept": "application/json",
+    }
+    data = {
+        "vector": query_vector,
+        "top_k": top_k,
+        "includeMetadata": include_metadata,
+        "includeValues": include_values,
+        "namespace": namespace,
+    }
+    return send_https_request(host, path, data, headers)
diff --git a/src/semantic_search/semantic_search/https_requests.py b/src/semantic_search/semantic_search/https_requests.py
new file mode 100644
index 0000000..edff750
--- /dev/null
+++ b/src/semantic_search/semantic_search/https_requests.py
@@ -0,0 +1,53 @@
+import json
+import socket
+import ssl
+
+
+def create_ssl_socket(host, port):
+    context = ssl.create_default_context()
+    sock = socket.create_connection((host, port))
+    return context.wrap_socket(sock, server_hostname=host)
+
+
+def parse_http_response(raw_response):
+    # Split response into headers and body
+    headers, _, body = raw_response.partition('\r\n\r\n')
+    # Parse the status line and headers
+    header_lines = headers.split('\r\n')
+    status_line = header_lines[0]
+    status_code = int(status_line.split()[1])  # Extract status code
+    header_dict = {k: v for k, v in (line.split(': ', 1) for line in header_lines[1:])}
+    return status_code, status_line, header_dict, body
+
+
+def read_all_from_socket(ssock):
+    buffer = b""
+    while True:
+        data = ssock.recv(4096)
+        if not data:
+            break
+        buffer += data
+    return buffer.decode()
+
+
+def send_https_request(host, path, data, headers=None):
+    if headers is None:
+        headers = {}
+    json_data = json.dumps(data)
+    content_length = len(json_data)
+    request_headers = {
+        "Host": host,
+        "Content-Type": "application/json",
+        "Content-Length": str(content_length),
+        "Connection": "close"
+    }
+    request_headers.update(headers)
+    request = f"POST {path} HTTP/1.1\r\n"
+    request += "".join(f"{key}: {value}\r\n" for key, value in request_headers.items())
+    request += "\r\n" + json_data
+    with create_ssl_socket(host, 443) as ssock:
+        ssock.send(request.encode())
+        raw_response = read_all_from_socket(ssock)
+
+        status_code, status_line, response_headers, response_body = parse_http_response(raw_response)
+        return status_code, status_line, response_headers, response_body
diff --git a/src/semantic_search/semantic_search/query.py b/src/semantic_search/semantic_search/query.py
index 32644a4..b2bfb5f 100644
--- a/src/semantic_search/semantic_search/query.py
+++ b/src/semantic_search/semantic_search/query.py
@@ -2,7 +2,7 @@
 import logging
 import time
 from datetime import date
-from .external_services.pinecone import get_pinecone_index
+from .external_services.pinecone import get_pinecone_index, query_index
 from .external_services.openai import create_embedding, gpt_query
 
 
@@ -36,15 +36,23 @@ def smart_query(namespace, query, username: str):
     logging.info(f"Smart Query: embedding created in {round(time.perf_counter() - stage_start_time, 2)}s")
 
     stage_start_time = time.perf_counter()
-    query_results = get_pinecone_index().query(
-        queries=[query_vector],
+    # query_results = get_pinecone_index().query(
+    #     queries=[query_vector],
+    #     top_k=50,
+    #     namespace=namespace,
+    #     include_values=False,
+    #     includeMetadata=True
+    # )
+    query_results = query_index(
+        query_vector=query_vector,
         top_k=50,
         namespace=namespace,
         include_values=False,
-        includeMetadata=True
+        include_metadata=True
     )
     logging.info(f"Smart Query: Pinecone search finished in {round(time.perf_counter() - stage_start_time, 2)}s")
-    query_matches = query_results['results'][0]['matches']
+    # query_matches = query_results['results'][0]['matches']
+    query_matches = json.loads(query_results[3])['matches']
 
     messages_for_gpt = [
         {