Weaviate Cookbooks

class QueryRewriter(dspy.Signature):
    """
    Rewrite the user's query into a more relevant search term that is a more relevant search term for searching a database.
    """
    input_query: str = dspy.InputField(description="The original user query")
    rewritten_query: str = dspy.OutputField(
        description=(
            "A single search term that is more relevant to the user's query. "
            "Include only relevant information, it does not need to be a full sentence or question "
        )
    )

def query_transformation(query: str) -> list[str]:
    lm = dspy.LM(subtask_model_name)
    answer = dspy.Predict(QueryRewriter)
    pred = answer(input_query=query, lm=lm)
    return [pred.rewritten_query]

class QueryRewriter(dspy.Signature):
    """
    Rewrite the user's query into a more relevant search terms that are more relevant search term for searching a database.
    """
    input_query: str = dspy.InputField(description="The original user query")
    rewritten_queries: list[str] = dspy.OutputField(
        description=(
            "A list of search terms that are more relevant to the user's query. "
            "Each entry should include only relevant information, it does not need to be a full sentence or question "
            "Split independent searches into different entries "
            "Each entry should be relevant independently that capture a different required search aspect "
            "Do not repeat similar search terms, each one should have a unique meaning "
            "Be sparse, do not duplicate search terms "
        )
    )

def query_transformation(query: str) -> list[str]:
    lm = dspy.LM(subtask_model_name)
    answer = dspy.Predict(QueryRewriter)
    pred = answer(input_query=query, lm=lm)
    return pred.rewritten_queries

from pydantic import BaseModel, Field
from typing import Literal, Any

class SearchFilter(BaseModel):
    field: str = Field(description="The field to be filtered on.")
    operator: Literal["=", "!=", ">", "<"] = Field(description="The operator to be used in conjunction with the value. These are strict operators.")
    value: Any = Field(description="The value to be used in conjunction with the operator.")

class Search(BaseModel):
    filters: list[SearchFilter] = Field(description="The filters to be used in the vector database. This is an AND operation.")

class SearchCreation(dspy.Signature):
    """
    Create filters and search parameters for a search query in a database.
    """
    query: str = dspy.InputField()
    schema: list[dict] = dspy.InputField(desc="Schema of the collection to be searched.")
    data_sample: list[dict] = dspy.InputField(desc="A sample of the data in the collection to be searched.")
    search: Search = dspy.OutputField(
        desc=(
            "Your filters and search parameters, this should be a valid JSON object. "
            "This should be constructed so that it matches the goal of the user prompt."
        )
    )

def _format_filters(search_filters: list[SearchFilter]):
    filters = []
    for search_filter in search_filters:
        base_filter = Filter.by_property(search_filter.field)
        if search_filter.operator == "=":
            filter = base_filter.equal(search_filter.value)
        elif search_filter.operator == "!=":
            filter = base_filter.not_equal(search_filter.value)
        elif search_filter.operator == ">":
            filter = base_filter.greater_than(search_filter.value)
        elif search_filter.operator == "<":
            filter = base_filter.less_than(search_filter.value)
        filters.append(filter)
    return Filter.all_of(filters) if filters else None

def create_filters(query: str):
    
    # import client here

    collection = client.collections.use("<collection_name>")

    # Get collection schema (for field names etc.). can replace this with more advanced configuration (like aggregating for unique groups)
    config = collection.config.get()
    schema = [{"name": p.name, "type": p.data_type[:]} for p in config.properties]

    # Get a sample of the data in the collection to be searched
    data_sample = collection.query.fetch_objects(limit=5)

    # Create search parameters
    search_parameters = dspy.ChainOfThought(SearchCreation)
    search_parameters_output = search_parameters(query=query, schema=schema, data_sample=data_sample, lm=dspy.LM(subtask_model_name))
    
    return _format_filters(search_parameters_output.search.filters)

collection = client.collections.use("<collection_name>")
collection.config.update(
    reranker_config=Reconfigure.Reranker.cohere()  
)

from weaviate.classes.query import Rerank

def retrieve(query: str, limit: int | None = None, filters = []) -> list[dict]:

    # ...existing code
    
    response = collection.query.hybrid(
        query=query,
        limit=limit,
        rerank=Rerank(
            prop="content", # what field to re-rank on
            query=query # what the search term for the re-ranker should be (same as original in this case)
        ),
        filters=filters if filters else None
    )
    
    # ...existing code

class Generator(dspy.Signature):
    """
    Answer the question based on the context.
    Do not include any information from external sources, only use the information provided in the context.
    If you cannot answer the question based on the information provided, say "I don't know".
    """
    context: str | list[dict] =  dspy.InputField(desc="The context to answer the question.")
    query: str = dspy.InputField(desc="The question to answer.")
    answer: str = dspy.OutputField(desc="The single answer to the question with no additional communication")

def generate(query: str, context: list[dict]) -> str:
    lm = dspy.LM(generation_model_name)
    answer = dspy.Predict(Generator)
    pred = answer(context=context, query=query, lm=lm)
    return pred.answer

from weaviate.agents.query import QueryAgent

# import client here

qa = QueryAgent(
    client=client, collections=["Example_Communications_Raw"]
)
response = qa.search("<user query here>") # just search with no text response
response = qa.ask("<user query here>") # search with text response accessible via response.final_answer