(self, state: dict)
| 41 | ) |
| 42 | |
| 43 | def execute(self, state: dict) -> dict: |
| 44 | self.logger.info(f"--- Executing {self.node_name} Node ---") |
| 45 | |
| 46 | try: |
| 47 | from qdrant_client import QdrantClient |
| 48 | from qdrant_client.models import Distance, PointStruct, VectorParams |
| 49 | except ImportError: |
| 50 | raise ImportError( |
| 51 | "qdrant_client is not installed. Please install it using 'pip install qdrant-client'." |
| 52 | ) |
| 53 | |
| 54 | if self.node_config.get("client_type") in ["memory", None]: |
| 55 | client = QdrantClient(":memory:") |
| 56 | elif self.node_config.get("client_type") == "local_db": |
| 57 | client = QdrantClient(path="path/to/db") |
| 58 | elif self.node_config.get("client_type") == "image": |
| 59 | client = QdrantClient(url="http://localhost:6333") |
| 60 | else: |
| 61 | raise ValueError("client_type provided not correct") |
| 62 | |
| 63 | docs = [elem.get("summary") for elem in state.get("docs")] |
| 64 | ids = list(range(1, len(state.get("docs")) + 1)) |
| 65 | |
| 66 | if state.get("embeddings"): |
| 67 | import openai |
| 68 | |
| 69 | openai_client = openai.Client() |
| 70 | |
| 71 | files = state.get("documents") |
| 72 | |
| 73 | array_of_embeddings = [] |
| 74 | i = 0 |
| 75 | |
| 76 | for file in files: |
| 77 | embeddings = openai_client.embeddings.create( |
| 78 | input=file, model=state.get("embeddings").get("model") |
| 79 | ) |
| 80 | i += 1 |
| 81 | points = PointStruct( |
| 82 | id=i, |
| 83 | vector=embeddings, |
| 84 | payload={"text": file}, |
| 85 | ) |
| 86 | |
| 87 | array_of_embeddings.append(points) |
| 88 | |
| 89 | collection_name = "collection" |
| 90 | |
| 91 | client.create_collection( |
| 92 | collection_name, |
| 93 | vectors_config=VectorParams( |
| 94 | size=1536, |
| 95 | distance=Distance.COSINE, |
| 96 | ), |
| 97 | ) |
| 98 | client.upsert(collection_name, points) |
| 99 | |
| 100 | state["vectorial_db"] = client |
no test coverage detected