使用 Qdrant 进行代码语义搜索

{
   "name":"await_ready_for_timeout",
   "signature":"fn await_ready_for_timeout (& self , timeout : Duration) -> bool",
   "code_type":"Function",
   "docstring":"= \" Return `true` if ready, `false` if timed out.\"",
   "line":44,
   "line_from":43,
   "line_to":51,
   "context":{
      "module":"common",
      "file_path":"lib/collection/src/common/is_ready.rs",
      "file_name":"is_ready.rs",
      "struct_name":"IsReady",
      "snippet":"    /// Return `true` if ready, `false` if timed out.\n    pub fn await_ready_for_timeout(&self, timeout: Duration) -> bool {\n        let mut is_ready = self.value.lock();\n        if !*is_ready {\n            !self.condvar.wait_for(&mut is_ready, timeout).timed_out()\n        } else {\n            true\n        }\n    }\n"
   }
}

wget https://storage.googleapis.com/tutorial-attachments/code-search/structures.jsonl

import json

structures = []
with open("structures.jsonl", "r") as fp:
    for i, row in enumerate(fp):
        entry = json.loads(row)
        structures.append(entry)

pip install inflection

import inflection
import re

from typing import Dict, Any

def textify(chunk: Dict[str, Any]) -> str:
    # Get rid of all the camel case / snake case
    # - inflection.underscore changes the camel case to snake case
    # - inflection.humanize converts the snake case to human readable form
    name = inflection.humanize(inflection.underscore(chunk["name"]))
    signature = inflection.humanize(inflection.underscore(chunk["signature"]))

    # Check if docstring is provided
    docstring = ""
    if chunk["docstring"]:
        docstring = f"that does {chunk['docstring']} "

    # Extract the location of that snippet of code
    context = (
        f"module {chunk['context']['module']} "
        f"file {chunk['context']['file_name']}"
    )
    if chunk["context"]["struct_name"]:
        struct_name = inflection.humanize(
            inflection.underscore(chunk["context"]["struct_name"])
        )
        context = f"defined in struct {struct_name} {context}"

    # Combine all the bits and pieces together
    text_representation = (
        f"{chunk['code_type']} {name} "
        f"{docstring}"
        f"defined as {signature} "
        f"{context}"
    )

    # Remove any special characters and concatenate the tokens
    tokens = re.split(r"\W", text_representation)
    tokens = filter(lambda x: x, tokens)
    return " ".join(tokens)

text_representations = list(map(textify, structures))

Function Await ready for timeout that does Return true if ready false if timed out defined as Fn await ready for timeout self timeout duration bool defined in struct Is ready module common file is_ready rs

pip install "qdrant-client[fastembed]"

QDRANT_URL = "https://my-cluster.cloud.qdrant.io:6333" # https://:6333 for local instance
QDRANT_API_KEY = "THIS_IS_YOUR_API_KEY" # None for local instance

from qdrant_client import QdrantClient, models

client = QdrantClient(QDRANT_URL, api_key=QDRANT_API_KEY)
client.create_collection(
    "qdrant-sources",
    vectors_config={
        "text": models.VectorParams(
            size=client.get_embedding_size(
                model_name="sentence-transformers/all-MiniLM-L6-v2"
            ),
            distance=models.Distance.COSINE,
        ),
        "code": models.VectorParams(
            size=client.get_embedding_size(
                model_name="jinaai/jina-embeddings-v2-base-code"
            ),
            distance=models.Distance.COSINE,
        ),
    },
)

import uuid

# Extract the code snippets from the structures to a separate list
code_snippets = [
    structure["context"]["snippet"] for structure in structures
]

points = [
    models.PointStruct(
        id=uuid.uuid4().hex,
        vector={
            "text": models.Document(
                text=text, model="sentence-transformers/all-MiniLM-L6-v2"
            ),
            "code": models.Document(
                text=code, model="jinaai/jina-embeddings-v2-base-code"
            ),
        },
        payload=structure,
    )
    for text, code, structure in zip(text_representations, code_snippets, structures)
]

# Note: This might take a while since inference happens implicitly.
# Parallel processing can help.
# But too many processes may trigger swap memory and hurt performance.
client.upload_points("qdrant-sources", points=points, batch_size=64)

query = "How do I count points in a collection?"

hits = client.query_points(
    "qdrant-sources",
    query=models.Document(text=query, model="sentence-transformers/all-MiniLM-L6-v2"),
    using="text",
    limit=5,
).points

hits = client.query_points(
    "qdrant-sources",
    query=models.Document(text=query, model="jinaai/jina-embeddings-v2-base-code"),
    using="code",
    limit=5,
).points

responses = client.query_batch_points(
    collection_name="qdrant-sources",
    requests=[
        models.QueryRequest(
            query=models.Document(
                text=query, model="sentence-transformers/all-MiniLM-L6-v2"
            ),
            using="text",
            with_payload=True,
            limit=5,
        ),
        models.QueryRequest(
            query=models.Document(
                text=query, model="jinaai/jina-embeddings-v2-base-code"
            ),
            using="code",
            with_payload=True,
            limit=5,
        ),
    ],
)

results = [response.points for response in responses]

results = client.query_points_groups(
    collection_name="qdrant-sources",
    using="code",
    query=models.Document(text=query, model="jinaai/jina-embeddings-v2-base-code"),
    group_by="context.module",
    limit=5,
    group_size=1,
)