Source code for langchain_azure_ai.query_constructors.cosmosdb_no_sql
"""Translator that converts a StructuredQuery into a CosmosDB NoSQL query."""fromtypingimportAny,Dict,Tuplefromlangchain_core.structured_queryimport(Comparator,Comparison,Operation,Operator,StructuredQuery,Visitor,)SQL_COMPARATOR={Comparator.EQ:"=",Comparator.NE:"!=",Comparator.GT:">",Comparator.GTE:">=",Comparator.LT:"<",Comparator.LTE:"<=",Comparator.LIKE:"LIKE",Comparator.IN:"IN",Comparator.NIN:"NOT IN",}SQL_OPERATOR={Operator.AND:"AND",Operator.OR:"OR",Operator.NOT:"NOT",}
[docs]classAzureCosmosDbNoSQLTranslator(Visitor):"""A visitor that converts a StructuredQuery into an CosmosDB NO SQL query."""
[docs]def__init__(self,table_name:str="c")->None:"""Initialize the translator with the table name."""self.table_name=table_name
[docs]defvisit_comparison(self,comparison:Comparison)->str:"""Visit a comparison operation and convert it into an SQL condition."""operator=SQL_COMPARATOR.get(comparison.comparator)value=comparison.valuefield=f"{self.table_name}.{comparison.attribute}"ifoperatorisNone:raiseValueError(f"Unsupported operator: {comparison.comparator}")# Correct value formattingifisinstance(value,str):value=f"'{value}'"elifisinstance(value,(list,tuple)):# Handle IN clauseifcomparison.comparatornotin[Comparator.IN,Comparator.NIN]:raiseValueError(f"Invalid comparator for list value: {comparison.comparator}")value=("("+", ".join(f"'{v}'"ifisinstance(v,str)elsestr(v)forvinvalue)+")")returnf"{field}{operator}{value}"
[docs]defvisit_operation(self,operation:Operation)->str:"""Visit logical operations and convert them into SQL expressions. Uses parentheses to ensure correct precedence. """operator=SQL_OPERATOR.get(operation.operator)ifoperatorisNone:raiseValueError(f"Unsupported operator: {operation.operator}")expressions=[arg.accept(self)forarginoperation.arguments]ifoperation.operator==Operator.NOT:returnf"NOT ({expressions[0]})"returnf"({f' {operator} '.join(expressions)})"
[docs]defvisit_structured_query(self,structured_query:StructuredQuery)->Tuple[str,Dict[str,Any]]:"""Visit a structured query and convert it into parameter for vectorstore."""ifstructured_query.filterisNone:kwargs={}else:kwargs={"where":structured_query.filter.accept(self)}returnstructured_query.query,kwargs