o
    Ai                     @   sJ   d dl Z d dlZd dlZd dlZd dlmZmZ h dZG dd dZ	dS )    N)ListDict>	   textdoc_idchunk_iddoc_name	tenant_id
source_urlchunk_orderpage_numbersection_headingc                   @   s   e Zd ZdZdefddZedejdejfddZ	ed	e
fd
dZdejdee
 fddZddejdefddZdefddZedefddZdS )
FaissStorezb
    Enterprise-grade FAISS vector store with
    strict metadata validation (POC compliant).
    dimc                 C   s   || _ t|| _g | _d S )N)r   faissIndexFlatIPindexmeta)selfr    r   E/var/www/html/fyndo/python/python_agents/rag_suite/rag/vectorstore.py__init__   s   
zFaissStore.__init__vectorsreturnc                 C   s   t jj| dddd }| | S )N   T)axiskeepdimsg-q=)nplinalgnorm)r   r   r   r   r   
_normalize"   s   zFaissStore._normalizer   c                 C   s&   t t|   }|rtd| d S )Nz"Missing required metadata fields: )REQUIRED_META_FIELDSsetkeys
ValueError)r   missingr   r   r   _validate_metadata(   s   zFaissStore._validate_metadata	metadatasc                 C   sX   t |t |krtd| |d}|D ]}| | q| j| | j| d S )Nz$Vectors and metadata length mismatchfloat32)	lenr#   r   astyper%   r   addr   extend)r   r   r&   mr   r   r   r*   3   s   zFaissStore.add   	query_veckc                 C   sj   |  |d}| j||\}}g }t|d |d D ]\}}|dkr%q|t|| j| d q|S )Nr'   r   )scoremetadata)r   r)   r   searchzipappendfloatr   )r   r.   r/   scoresindicesresultsr1   idxr   r   r   r3   B   s   
zFaissStore.searchpathc                 C   s   t j|dd t| jt j|d tt j|dddd}tj	| j
t| j| jd|d	d
d W d    d S 1 s=w   Y  d S )NT)exist_okindex.faiss	meta.jsonwutf-8encoding)r   countr      F)indentensure_ascii)osmakedirsr   write_indexr   r;   joinopenjsondumpr   r(   r   )r   r;   fr   r   r   saveS   s    "zFaissStore.savec                 C   s   t j|d}t j|d}t j|rt j|stdt|}t|ddd}t	|}W d    n1 s:w   Y  | |d d}||_
|d	 |_|S )
Nr=   r>   z!FAISS index or metadata not foundrr@   rA   r   )r   r   )rG   r;   rJ   existsFileNotFoundErrorr   
read_indexrK   rL   loadr   r   )clsr;   
index_path	meta_pathr   rN   objstorer   r   r   rT   h   s   

zFaissStore.loadN)r-   )__name__
__module____qualname____doc__intr   staticmethodr   ndarrayr   r   r%   r   r*   r3   strrO   classmethodrT   r   r   r   r   r      s    
r   )
rG   rL   r   numpyr   typingr   r   r    r   r   r   r   r   <module>   s   