diff --git a/cognee-frontend/src/modules/chat/getHistory.ts b/cognee-frontend/src/modules/chat/getHistory.ts new file mode 100644 index 00000000..dce914da --- /dev/null +++ b/cognee-frontend/src/modules/chat/getHistory.ts @@ -0,0 +1,8 @@ +import { fetch } from '@/utils'; + +export default function getHistory() { + return fetch( + '/v1/search', + ) + .then((response) => response.json()); +} diff --git a/cognee-frontend/src/ui/Partials/SearchView/SearchView.tsx b/cognee-frontend/src/ui/Partials/SearchView/SearchView.tsx index b20beb5b..b4fa0777 100644 --- a/cognee-frontend/src/ui/Partials/SearchView/SearchView.tsx +++ b/cognee-frontend/src/ui/Partials/SearchView/SearchView.tsx @@ -1,9 +1,12 @@ +'use client'; + import { v4 } from 'uuid'; import classNames from 'classnames'; -import { useCallback, useState } from 'react'; +import { useCallback, useEffect, useState } from 'react'; import { CTAButton, Stack, Text, DropdownSelect, TextArea, useBoolean } from 'ohmy-ui'; import { fetch } from '@/utils'; import styles from './SearchView.module.css'; +import getHistory from '@/modules/chat/getHistory'; interface Message { id: string; @@ -52,6 +55,14 @@ export default function SearchView() { }, 300); }, []); + useEffect(() => { + getHistory() + .then((history) => { + setMessages(history); + scrollToBottom(); + }); + }, [scrollToBottom]); + const handleSearchSubmit = useCallback((event: React.FormEvent) => { event.preventDefault(); @@ -78,7 +89,7 @@ export default function SearchView() { 'Content-Type': 'application/json', }, body: JSON.stringify({ - query: inputValue, + query: inputValue.trim(), searchType: searchTypeValue, }), }) diff --git a/cognee/__init__.py b/cognee/__init__.py index aca7f5d4..e89ef1dc 100644 --- a/cognee/__init__.py +++ b/cognee/__init__.py @@ -2,7 +2,7 @@ from .api.v1.add import add from .api.v1.cognify import cognify from .api.v1.datasets.datasets import datasets -from .api.v1.search import search, SearchType +from .api.v1.search import search, SearchType, get_search_history from .api.v1.prune import prune # Pipelines diff --git a/cognee/api/v1/search/__init__.py b/cognee/api/v1/search/__init__.py index f01dcd63..91cf35c8 100644 --- a/cognee/api/v1/search/__init__.py +++ b/cognee/api/v1/search/__init__.py @@ -1 +1,2 @@ from .search_v2 import search, SearchType +from .get_search_history import get_search_history diff --git a/cognee/api/v1/search/get_search_history.py b/cognee/api/v1/search/get_search_history.py new file mode 100644 index 00000000..fada67c8 --- /dev/null +++ b/cognee/api/v1/search/get_search_history.py @@ -0,0 +1,9 @@ +from cognee.modules.search.operations import get_history +from cognee.modules.users.methods import get_default_user +from cognee.modules.users.models import User + +async def get_search_history(user: User = None) -> list: + if not user: + user = await get_default_user() + + return await get_history(user.id) diff --git a/cognee/api/v1/search/routers/get_search_router.py b/cognee/api/v1/search/routers/get_search_router.py index 5df49635..893067c2 100644 --- a/cognee/api/v1/search/routers/get_search_router.py +++ b/cognee/api/v1/search/routers/get_search_router.py @@ -1,8 +1,11 @@ -from cognee.api.v1.search import SearchType +from uuid import UUID +from datetime import datetime +from fastapi import Depends, APIRouter from fastapi.responses import JSONResponse +from cognee.api.v1.search import SearchType +from cognee.api.DTO import InDTO, OutDTO from cognee.modules.users.models import User -from fastapi import Depends, APIRouter -from cognee.api.DTO import InDTO +from cognee.modules.search.operations import get_history from cognee.modules.users.methods import get_authenticated_user @@ -13,6 +16,24 @@ class SearchPayloadDTO(InDTO): def get_search_router() -> APIRouter: router = APIRouter() + class SearchHistoryItem(OutDTO): + id: UUID + text: str + user: str + created_at: datetime + + @router.get("/", response_model = list[SearchHistoryItem]) + async def get_search_history(user: User = Depends(get_authenticated_user)): + try: + history = await get_history(user.id) + + return history + except Exception as error: + return JSONResponse( + status_code = 500, + content = {"error": str(error)} + ) + @router.post("/", response_model = list) async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticated_user)): """ This endpoint is responsible for searching for nodes in the graph.""" @@ -28,4 +49,4 @@ async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticat content = {"error": str(error)} ) - return router \ No newline at end of file + return router diff --git a/cognee/api/v1/search/search_v2.py b/cognee/api/v1/search/search_v2.py index a82f1421..c1bc0ee4 100644 --- a/cognee/api/v1/search/search_v2.py +++ b/cognee/api/v1/search/search_v2.py @@ -1,6 +1,9 @@ +import json from uuid import UUID from enum import Enum from typing import Callable, Dict +from cognee.modules.search.operations import log_query, log_result +from cognee.modules.storage.utils import JSONEncoder from cognee.shared.utils import send_telemetry from cognee.modules.users.models import User from cognee.modules.users.methods import get_default_user @@ -14,15 +17,17 @@ class SearchType(Enum): INSIGHTS = "INSIGHTS" CHUNKS = "CHUNKS" -async def search(search_type: SearchType, query: str, user: User = None) -> list: +async def search(query_type: SearchType, query_text: str, user: User = None) -> list: if user is None: user = await get_default_user() if user is None: raise PermissionError("No user found in the system. Please create a user.") + query = await log_query(query_text, str(query_type), user.id) + own_document_ids = await get_document_ids_for_user(user.id) - search_results = await specific_search(search_type, query, user) + search_results = await specific_search(query_type, query_text, user) filtered_search_results = [] @@ -33,19 +38,21 @@ async def search(search_type: SearchType, query: str, user: User = None) -> list if document_id is None or document_id in own_document_ids: filtered_search_results.append(search_result) + await log_result(query.id, json.dumps(filtered_search_results, cls = JSONEncoder), user.id) + return filtered_search_results -async def specific_search(search_type: SearchType, query: str, user) -> list: +async def specific_search(query_type: SearchType, query: str, user) -> list: search_tasks: Dict[SearchType, Callable] = { SearchType.SUMMARIES: query_summaries, SearchType.INSIGHTS: query_graph_connections, SearchType.CHUNKS: query_chunks, } - search_task = search_tasks.get(search_type) + search_task = search_tasks.get(query_type) if search_task is None: - raise ValueError(f"Unsupported search type: {search_type}") + raise ValueError(f"Unsupported search type: {query_type}") send_telemetry("cognee.search EXECUTION STARTED", user.id) diff --git a/cognee/infrastructure/databases/relational/__init__.py b/cognee/infrastructure/databases/relational/__init__.py index 1ef84790..09a4d669 100644 --- a/cognee/infrastructure/databases/relational/__init__.py +++ b/cognee/infrastructure/databases/relational/__init__.py @@ -2,6 +2,3 @@ from .config import get_relational_config from .create_db_and_tables import create_db_and_tables from .get_relational_engine import get_relational_engine - -# Global data types -from .data_types.UUID import UUID diff --git a/cognee/infrastructure/databases/relational/data_types/UUID.py b/cognee/infrastructure/databases/relational/data_types/UUID.py deleted file mode 100644 index 722204b3..00000000 --- a/cognee/infrastructure/databases/relational/data_types/UUID.py +++ /dev/null @@ -1,45 +0,0 @@ -import uuid - -from sqlalchemy.types import TypeDecorator, BINARY -from sqlalchemy.dialects.postgresql import UUID as psqlUUID - -class UUID(TypeDecorator): - """Platform-independent GUID type. - - Uses Postgresql's UUID type, otherwise uses - BINARY(16), to store UUID. - - """ - impl = BINARY - - def load_dialect_impl(self, dialect): - if dialect.name == 'postgresql': - return dialect.type_descriptor(psqlUUID()) - else: - return dialect.type_descriptor(BINARY(16)) - - def process_bind_param(self, value, dialect): - if value is None: - return value - else: - if not isinstance(value, uuid.UUID): - if isinstance(value, bytes): - value = uuid.UUID(bytes = value) - elif isinstance(value, int): - value = uuid.UUID(int = value) - elif isinstance(value, str): - value = uuid.UUID(value) - if dialect.name == 'postgresql': - return str(value) - else: - return value.bytes - - def process_result_value(self, value, dialect): - if value is None: - return value - if dialect.name == 'postgresql': - if isinstance(value, uuid.UUID): - return value - return uuid.UUID(value) - else: - return uuid.UUID(bytes = value) diff --git a/cognee/infrastructure/pipeline/models/Operation.py b/cognee/infrastructure/pipeline/models/Operation.py index 1834c1a3..62eb74c4 100644 --- a/cognee/infrastructure/pipeline/models/Operation.py +++ b/cognee/infrastructure/pipeline/models/Operation.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import datetime, timezone from sqlalchemy.orm import Mapped, MappedColumn from sqlalchemy import Column, DateTime, ForeignKey, Enum, JSON from cognee.infrastructure.databases.relational import Base, UUID @@ -24,4 +24,4 @@ class Operation(Base): data_id = Column(UUID, ForeignKey("data.id")) meta_data: Mapped[dict] = MappedColumn(type_ = JSON) - created_at = Column(DateTime, default = datetime.utcnow) + created_at = Column(DateTime, default = datetime.now(timezone.utc)) diff --git a/cognee/modules/data/models/Data.py b/cognee/modules/data/models/Data.py index 06452153..2e974560 100644 --- a/cognee/modules/data/models/Data.py +++ b/cognee/modules/data/models/Data.py @@ -2,8 +2,8 @@ from typing import List from datetime import datetime, timezone from sqlalchemy.orm import relationship, Mapped -from sqlalchemy import Column, String, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, String, DateTime, UUID +from cognee.infrastructure.databases.relational import Base from .DatasetData import DatasetData class Data(Base): diff --git a/cognee/modules/data/models/Dataset.py b/cognee/modules/data/models/Dataset.py index 5cf5d235..f7078b8f 100644 --- a/cognee/modules/data/models/Dataset.py +++ b/cognee/modules/data/models/Dataset.py @@ -2,8 +2,8 @@ from typing import List from datetime import datetime, timezone from sqlalchemy.orm import relationship, Mapped -from sqlalchemy import Column, Text, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, Text, DateTime, UUID +from cognee.infrastructure.databases.relational import Base from .DatasetData import DatasetData class Dataset(Base): diff --git a/cognee/modules/data/models/DatasetData.py b/cognee/modules/data/models/DatasetData.py index ed9d3c64..a35c120e 100644 --- a/cognee/modules/data/models/DatasetData.py +++ b/cognee/modules/data/models/DatasetData.py @@ -1,6 +1,6 @@ from datetime import datetime, timezone -from sqlalchemy import Column, DateTime, ForeignKey -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, DateTime, ForeignKey, UUID +from cognee.infrastructure.databases.relational import Base class DatasetData(Base): __tablename__ = "dataset_data" diff --git a/cognee/modules/pipelines/models/Pipeline.py b/cognee/modules/pipelines/models/Pipeline.py index e9cad594..f4d20bb0 100644 --- a/cognee/modules/pipelines/models/Pipeline.py +++ b/cognee/modules/pipelines/models/Pipeline.py @@ -1,9 +1,10 @@ from uuid import uuid4 from datetime import datetime, timezone -from sqlalchemy import Column, DateTime, String, Text +from sqlalchemy import Column, DateTime, String, Text, UUID from sqlalchemy.orm import relationship, Mapped -from cognee.infrastructure.databases.relational import Base, UUID +from cognee.infrastructure.databases.relational import Base from .PipelineTask import PipelineTask +from .Task import Task class Pipeline(Base): __tablename__ = "pipelines" diff --git a/cognee/modules/pipelines/models/PipelineRun.py b/cognee/modules/pipelines/models/PipelineRun.py index 5d5969b2..ab3498ef 100644 --- a/cognee/modules/pipelines/models/PipelineRun.py +++ b/cognee/modules/pipelines/models/PipelineRun.py @@ -1,8 +1,8 @@ import enum from uuid import uuid4 from datetime import datetime, timezone -from sqlalchemy import Column, DateTime, JSON, Enum -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, DateTime, JSON, Enum, UUID +from cognee.infrastructure.databases.relational import Base class PipelineRunStatus(enum.Enum): DATASET_PROCESSING_STARTED = "DATASET_PROCESSING_STARTED" diff --git a/cognee/modules/pipelines/models/PipelineTask.py b/cognee/modules/pipelines/models/PipelineTask.py index acbf44e5..c6c7eb5e 100644 --- a/cognee/modules/pipelines/models/PipelineTask.py +++ b/cognee/modules/pipelines/models/PipelineTask.py @@ -1,6 +1,6 @@ from datetime import datetime, timezone -from sqlalchemy import Column, DateTime, ForeignKey -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, DateTime, ForeignKey, UUID +from cognee.infrastructure.databases.relational import Base class PipelineTask(Base): __tablename__ = "pipeline_task" diff --git a/cognee/modules/search/models/Query.py b/cognee/modules/search/models/Query.py new file mode 100644 index 00000000..18219633 --- /dev/null +++ b/cognee/modules/search/models/Query.py @@ -0,0 +1,16 @@ +from uuid import uuid4 +from datetime import datetime, timezone +from sqlalchemy import Column, DateTime, String, UUID +from cognee.infrastructure.databases.relational import Base + +class Query(Base): + __tablename__ = "queries" + + id = Column(UUID, primary_key = True, default = uuid4) + + text = Column(String) + query_type = Column(String) + user_id = Column(UUID) + + created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc)) + updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc)) diff --git a/cognee/modules/search/models/Result.py b/cognee/modules/search/models/Result.py new file mode 100644 index 00000000..acda59dd --- /dev/null +++ b/cognee/modules/search/models/Result.py @@ -0,0 +1,16 @@ +from datetime import datetime, timezone +from uuid import uuid4 +from sqlalchemy import Column, DateTime, Text, UUID +from cognee.infrastructure.databases.relational import Base + +class Result(Base): + __tablename__ = "results" + + id = Column(UUID, primary_key = True, default = uuid4) + + value = Column(Text) + query_id = Column(UUID) + user_id = Column(UUID, index = True) + + created_at = Column(DateTime(timezone = True), default = lambda: datetime.now(timezone.utc)) + updated_at = Column(DateTime(timezone = True), onupdate = lambda: datetime.now(timezone.utc)) diff --git a/cognee/modules/search/operations/__init__.py b/cognee/modules/search/operations/__init__.py new file mode 100644 index 00000000..41d2a4e4 --- /dev/null +++ b/cognee/modules/search/operations/__init__.py @@ -0,0 +1,3 @@ +from .log_query import log_query +from .log_result import log_result +from .get_history import get_history diff --git a/cognee/modules/search/operations/get_history.py b/cognee/modules/search/operations/get_history.py new file mode 100644 index 00000000..831c4acc --- /dev/null +++ b/cognee/modules/search/operations/get_history.py @@ -0,0 +1,31 @@ +from uuid import UUID +from sqlalchemy import literal, select +from cognee.infrastructure.databases.relational import get_relational_engine +from ..models.Query import Query +from ..models.Result import Result + +async def get_history(user_id: UUID, limit: int = 10) -> list[Result]: + db_engine = get_relational_engine() + + queries_query = select( + Query.id, + Query.text.label("text"), + Query.created_at, + literal("user").label("user") + ) \ + .filter(Query.user_id == user_id) + + results_query = select( + Result.id, + Result.value.label("text"), + Result.created_at, + literal("system").label("user") + ) \ + .filter(Result.user_id == user_id) + + history_query = queries_query.union(results_query).order_by("created_at").limit(limit) + + async with db_engine.get_async_session() as session: + history = (await session.execute(history_query)).all() + + return history diff --git a/cognee/modules/search/operations/get_queries.py b/cognee/modules/search/operations/get_queries.py new file mode 100644 index 00000000..ded10a8e --- /dev/null +++ b/cognee/modules/search/operations/get_queries.py @@ -0,0 +1,17 @@ +from uuid import UUID +from sqlalchemy import select +from cognee.infrastructure.databases.relational import get_relational_engine +from ..models.Query import Query + +async def get_queries(user_id: UUID, limit: int) -> list[Query]: + db_engine = get_relational_engine() + + async with db_engine.get_async_session() as session: + queries = (await session.scalars( + select(Query) + .filter(Query.user_id == user_id) + .order_by(Query.created_at.desc()) + .limit(limit) + )).all() + + return queries diff --git a/cognee/modules/search/operations/get_results.py b/cognee/modules/search/operations/get_results.py new file mode 100644 index 00000000..7f90a3f0 --- /dev/null +++ b/cognee/modules/search/operations/get_results.py @@ -0,0 +1,17 @@ +from uuid import UUID +from sqlalchemy import select +from cognee.infrastructure.databases.relational import get_relational_engine +from ..models.Result import Result + +async def get_results(user_id: UUID, limit: int = 10) -> list[Result]: + db_engine = get_relational_engine() + + async with db_engine.get_async_session() as session: + results = (await session.scalars( + select(Result) + .filter(Result.user_id == user_id) + .order_by(Result.created_at.desc()) + .limit(limit) + )).all() + + return results diff --git a/cognee/modules/search/operations/log_query.py b/cognee/modules/search/operations/log_query.py new file mode 100644 index 00000000..02ed3f15 --- /dev/null +++ b/cognee/modules/search/operations/log_query.py @@ -0,0 +1,19 @@ +from uuid import UUID +from cognee.infrastructure.databases.relational import get_relational_engine +from ..models.Query import Query + +async def log_query(query_text: str, query_type: str, user_id: UUID) -> Query: + db_engine = get_relational_engine() + + async with db_engine.get_async_session() as session: + query = Query( + text = query_text, + query_type = query_type, + user_id = user_id, + ) + + session.add(query) + + await session.commit() + + return query diff --git a/cognee/modules/search/operations/log_result.py b/cognee/modules/search/operations/log_result.py new file mode 100644 index 00000000..b81e0b44 --- /dev/null +++ b/cognee/modules/search/operations/log_result.py @@ -0,0 +1,15 @@ +from uuid import UUID +from cognee.infrastructure.databases.relational import get_relational_engine +from ..models.Result import Result + +async def log_result(query_id: UUID, result: str, user_id: UUID): + db_engine = get_relational_engine() + + async with db_engine.get_async_session() as session: + session.add(Result( + value = result, + query_id = query_id, + user_id = user_id, + )) + + await session.commit() diff --git a/cognee/modules/users/models/ACL.py b/cognee/modules/users/models/ACL.py index b01fe601..f54d2422 100644 --- a/cognee/modules/users/models/ACL.py +++ b/cognee/modules/users/models/ACL.py @@ -1,8 +1,8 @@ from uuid import uuid4 from datetime import datetime, timezone from sqlalchemy.orm import relationship, Mapped -from sqlalchemy import Column, ForeignKey, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, ForeignKey, DateTime, UUID +from cognee.infrastructure.databases.relational import Base from .ACLResources import ACLResources class ACL(Base): diff --git a/cognee/modules/users/models/ACLResources.py b/cognee/modules/users/models/ACLResources.py index 268d4a75..464fed2e 100644 --- a/cognee/modules/users/models/ACLResources.py +++ b/cognee/modules/users/models/ACLResources.py @@ -1,6 +1,6 @@ from datetime import datetime, timezone -from sqlalchemy import Column, ForeignKey, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, ForeignKey, DateTime, UUID +from cognee.infrastructure.databases.relational import Base class ACLResources(Base): __tablename__ = "acl_resources" diff --git a/cognee/modules/users/models/Group.py b/cognee/modules/users/models/Group.py index d86dbee9..793decb3 100644 --- a/cognee/modules/users/models/Group.py +++ b/cognee/modules/users/models/Group.py @@ -1,6 +1,5 @@ from sqlalchemy.orm import relationship, Mapped -from sqlalchemy import Column, String, ForeignKey -from cognee.infrastructure.databases.relational import UUID +from sqlalchemy import Column, String, ForeignKey, UUID from .Principal import Principal from .UserGroup import UserGroup diff --git a/cognee/modules/users/models/Permission.py b/cognee/modules/users/models/Permission.py index 84b1a307..3b170937 100644 --- a/cognee/modules/users/models/Permission.py +++ b/cognee/modules/users/models/Permission.py @@ -1,8 +1,8 @@ from uuid import uuid4 from datetime import datetime, timezone # from sqlalchemy.orm import relationship -from sqlalchemy import Column, DateTime, String -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, DateTime, String, UUID +from cognee.infrastructure.databases.relational import Base class Permission(Base): __tablename__ = "permissions" diff --git a/cognee/modules/users/models/Principal.py b/cognee/modules/users/models/Principal.py index 4ef91ffa..dc6e5130 100644 --- a/cognee/modules/users/models/Principal.py +++ b/cognee/modules/users/models/Principal.py @@ -1,7 +1,7 @@ from uuid import uuid4 from datetime import datetime, timezone -from sqlalchemy import Column, String, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, String, DateTime, UUID +from cognee.infrastructure.databases.relational import Base class Principal(Base): __tablename__ = "principals" diff --git a/cognee/modules/users/models/Resource.py b/cognee/modules/users/models/Resource.py index 0eca509a..563f9627 100644 --- a/cognee/modules/users/models/Resource.py +++ b/cognee/modules/users/models/Resource.py @@ -1,8 +1,8 @@ from uuid import uuid4 from datetime import datetime, timezone from sqlalchemy.orm import relationship -from sqlalchemy import Column, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, DateTime, UUID +from cognee.infrastructure.databases.relational import Base from .ACLResources import ACLResources class Resource(Base): diff --git a/cognee/modules/users/models/User.py b/cognee/modules/users/models/User.py index 96f78b12..3536ac94 100644 --- a/cognee/modules/users/models/User.py +++ b/cognee/modules/users/models/User.py @@ -1,10 +1,10 @@ from uuid import UUID as uuid_UUID -from sqlalchemy import ForeignKey, Column +from sqlalchemy import ForeignKey, Column, UUID from sqlalchemy.orm import relationship, Mapped from fastapi_users.db import SQLAlchemyBaseUserTableUUID -from cognee.infrastructure.databases.relational import UUID from .Principal import Principal from .UserGroup import UserGroup +from .Group import Group class User(SQLAlchemyBaseUserTableUUID, Principal): __tablename__ = "users" @@ -25,7 +25,6 @@ class User(SQLAlchemyBaseUserTableUUID, Principal): from fastapi_users import schemas class UserRead(schemas.BaseUser[uuid_UUID]): - # groups: list[uuid_UUID] # Add groups attribute pass class UserCreate(schemas.BaseUserCreate): diff --git a/cognee/modules/users/models/UserGroup.py b/cognee/modules/users/models/UserGroup.py index a2dfa8bc..5a85c9d3 100644 --- a/cognee/modules/users/models/UserGroup.py +++ b/cognee/modules/users/models/UserGroup.py @@ -1,6 +1,6 @@ from datetime import datetime, timezone -from sqlalchemy import Column, ForeignKey, DateTime -from cognee.infrastructure.databases.relational import Base, UUID +from sqlalchemy import Column, ForeignKey, DateTime, UUID +from cognee.infrastructure.databases.relational import Base class UserGroup(Base): __tablename__ = "user_groups" diff --git a/cognee/tests/integration/run_toy_tasks/data/cognee_db b/cognee/tests/integration/run_toy_tasks/data/cognee_db index 60455ad2..912adf29 100644 Binary files a/cognee/tests/integration/run_toy_tasks/data/cognee_db and b/cognee/tests/integration/run_toy_tasks/data/cognee_db differ diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py index 2e707b64..88c9cdc7 100755 --- a/cognee/tests/test_library.py +++ b/cognee/tests/test_library.py @@ -35,24 +35,27 @@ async def main(): random_node = (await vector_engine.search("Entity_name", "AI"))[0] random_node_name = random_node.payload["text"] - search_results = await cognee.search(SearchType.INSIGHTS, query = random_node_name) + search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted sentences are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.CHUNKS, query = random_node_name) + search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted chunks are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name) + search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name) assert len(search_results) != 0, "Query related summaries don't exist." print("\nExtracted summaries are:\n") for result in search_results: print(f"{result}\n") + history = await cognee.get_search_history() + + assert len(history) == 6, "Search history is not correct." if __name__ == "__main__": import asyncio diff --git a/cognee/tests/test_neo4j.py b/cognee/tests/test_neo4j.py index 9cf1c53d..44bacd9b 100644 --- a/cognee/tests/test_neo4j.py +++ b/cognee/tests/test_neo4j.py @@ -39,24 +39,27 @@ async def main(): random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] random_node_name = random_node.payload["text"] - search_results = await cognee.search(SearchType.INSIGHTS, query = random_node_name) + search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted sentences are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.CHUNKS, query = random_node_name) + search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted chunks are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name) + search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name) assert len(search_results) != 0, "Query related summaries don't exist." print("\nExtracted summaries are:\n") for result in search_results: print(f"{result}\n") + history = await cognee.get_search_history() + + assert len(history) == 6, "Search history is not correct." if __name__ == "__main__": import asyncio diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index ac4d08fb..b5a6fc44 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -68,24 +68,27 @@ async def main(): random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] random_node_name = random_node.payload["text"] - search_results = await cognee.search(SearchType.INSIGHTS, query=random_node_name) + search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted sentences are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.CHUNKS, query=random_node_name) + search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted chunks are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.SUMMARIES, query=random_node_name) + search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name) assert len(search_results) != 0, "Query related summaries don't exist." print("\n\nExtracted summaries are:\n") for result in search_results: print(f"{result}\n") + history = await cognee.get_search_history() + + assert len(history) == 6, "Search history is not correct." if __name__ == "__main__": import asyncio diff --git a/cognee/tests/test_qdrant.py b/cognee/tests/test_qdrant.py index 784b3f27..8ca525f0 100644 --- a/cognee/tests/test_qdrant.py +++ b/cognee/tests/test_qdrant.py @@ -40,24 +40,27 @@ async def main(): random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] random_node_name = random_node.payload["text"] - search_results = await cognee.search(SearchType.INSIGHTS, query = random_node_name) + search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted sentences are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.CHUNKS, query = random_node_name) + search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted chunks are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name) + search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name) assert len(search_results) != 0, "Query related summaries don't exist." print("\nExtracted summaries are:\n") for result in search_results: print(f"{result}\n") + history = await cognee.get_search_history() + + assert len(history) == 6, "Search history is not correct." if __name__ == "__main__": import asyncio diff --git a/cognee/tests/test_weaviate.py b/cognee/tests/test_weaviate.py index f788f997..cdb4b934 100644 --- a/cognee/tests/test_weaviate.py +++ b/cognee/tests/test_weaviate.py @@ -38,24 +38,27 @@ async def main(): random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0] random_node_name = random_node.payload["text"] - search_results = await cognee.search(SearchType.INSIGHTS, query = random_node_name) + search_results = await cognee.search(SearchType.INSIGHTS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted sentences are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.CHUNKS, query = random_node_name) + search_results = await cognee.search(SearchType.CHUNKS, query_text = random_node_name) assert len(search_results) != 0, "The search results list is empty." print("\n\nExtracted chunks are:\n") for result in search_results: print(f"{result}\n") - search_results = await cognee.search(SearchType.SUMMARIES, query = random_node_name) + search_results = await cognee.search(SearchType.SUMMARIES, query_text = random_node_name) assert len(search_results) != 0, "Query related summaries don't exist." print("\nExtracted summaries are:\n") for result in search_results: print(f"{result}\n") + history = await cognee.get_search_history() + + assert len(history) == 6, "Search history is not correct." if __name__ == "__main__": import asyncio diff --git a/notebooks/cognee_demo.ipynb b/notebooks/cognee_demo.ipynb index 06cd2a86..45f5a618 100644 --- a/notebooks/cognee_demo.ipynb +++ b/notebooks/cognee_demo.ipynb @@ -791,7 +791,7 @@ "node = (await vector_engine.search(\"Entity_name\", \"sarah.nguyen@example.com\"))[0]\n", "node_name = node.payload[\"text\"]\n", "\n", - "search_results = await cognee.search(SearchType.SUMMARIES, query = node_name)\n", + "search_results = await cognee.search(SearchType.SUMMARIES, query_text = node_name)\n", "print(\"\\n\\Extracted summaries are:\\n\")\n", "for result in search_results:\n", " print(f\"{result}\\n\")" @@ -812,7 +812,7 @@ "metadata": {}, "outputs": [], "source": [ - "search_results = await cognee.search(SearchType.CHUNKS, query = node_name)\n", + "search_results = await cognee.search(SearchType.CHUNKS, query_text = node_name)\n", "print(\"\\n\\nExtracted chunks are:\\n\")\n", "for result in search_results:\n", " print(f\"{result}\\n\")" @@ -833,7 +833,7 @@ "metadata": {}, "outputs": [], "source": [ - "search_results = await cognee.search(SearchType.INSIGHTS, query = node_name)\n", + "search_results = await cognee.search(SearchType.INSIGHTS, query_text = node_name)\n", "print(\"\\n\\nExtracted sentences are:\\n\")\n", "for result in search_results:\n", " print(f\"{result}\\n\")"