o
    Zh                     @   sH   d dl mZmZmZ d dlmZmZmZmZm	Z	m
Z
 G dd de
ZdS )    )DictTupleUnion)
Comparator
Comparison	OperationOperatorStructuredQueryVisitorc                   @   s   e Zd ZdZejejejejej	ej
ejgZ	 ejejejgZ	 deeef defddZdedefddZd	edefd
dZdedeeef fddZdS )ElasticsearchTranslatorzLTranslate `Elasticsearch` internal query language elements to valid filters.funcreturnc                 C   sR   |  | tjdtjdtjdtjdtjdtjdtj	dtj
dtjd	tjd	i
}|| S )
NZshouldZmust_notZmusttermgtZgteltZltematch)Z_validate_funcr   ORNOTANDr   EQGTGTELTLTECONTAINLIKE)selfr   Zmap_dict r   k/var/www/html/lang_env/lib/python3.10/site-packages/langchain_community/query_constructors/elasticsearch.py_format_func   s   
z$ElasticsearchTranslator._format_func	operationc                    s(    fdd|j D }d |j|iiS )Nc                    s   g | ]}|  qS r   )accept).0argr   r   r   
<listcomp>/   s    z;ElasticsearchTranslator.visit_operation.<locals>.<listcomp>bool)	argumentsr   operator)r   r    argsr   r$   r   visit_operation.   s   z'ElasticsearchTranslator.visit_operation
comparisonc                 C   s  d|j  }|jtjtjtjtjfv }|r4|j}t|jt	r(d|jv r(|jd }d|| 
|j|iiiS |jtjkrG| 
|j|d|jiiiS |jtjkr[| 
|j||jddiiS t|jtrf| dn|}t|jt	ryd|jv ry|jd |_| 
|j||jiiS )Nz	metadata.daterangequeryZAUTO)r.   Z	fuzzinessz.keyword)	attributeZ
comparatorr   r   r   r   r   value
isinstancedictr   r   r   str)r   r+   fieldZis_range_comparatorr0   r   r   r   visit_comparison3   s6   




z(ElasticsearchTranslator.visit_comparisonstructured_queryc                 C   s,   |j d u ri }n	d|j | gi}|j|fS )Nfilter)r7   r!   r.   )r   r6   kwargsr   r   r   visit_structured_query]   s   

z.ElasticsearchTranslator.visit_structured_queryN)__name__
__module____qualname____doc__r   r   r   r   r   r   r   r   Zallowed_comparatorsr   r   r   r   Zallowed_operatorsr   r3   r   r   r   r*   r   r5   r	   r   r2   r9   r   r   r   r   r      s*    	*
r   N)typingr   r   r   Zlangchain_core.structured_queryr   r   r   r   r	   r
   r   r   r   r   r   <module>   s     
