-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Cog 334 structure routing #164
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
f09c28a
f8f2746
4fd4651
02d6750
fc3eb9b
eca1b9f
de1ba5c
d5a220e
801efeb
7425409
bce6540
ddf4952
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
Resolved security concerns in endpoints regarding os Fix #COG-334-structure-routing
- Loading branch information
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -13,20 +13,30 @@ | |||||||||||||||||||||||||||||||||||||||||||||||||
| from cognee.modules.users.methods import get_authenticated_user | ||||||||||||||||||||||||||||||||||||||||||||||||||
| from cognee.modules.pipelines.models import PipelineRunStatus | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| def get_datasets_router(): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| logger = logging.getLogger(__name__) | ||||||||||||||||||||||||||||||||||||||||||||||||||
| logger = logging.getLogger(__name__) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| class ErrorResponseDTO(BaseModel): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| message: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| class DatasetDTO(OutDTO): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| id: UUID | ||||||||||||||||||||||||||||||||||||||||||||||||||
| name: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| created_at: datetime | ||||||||||||||||||||||||||||||||||||||||||||||||||
| updated_at: Optional[datetime] | ||||||||||||||||||||||||||||||||||||||||||||||||||
| owner_id: UUID | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| class DataDTO(OutDTO): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| id: UUID | ||||||||||||||||||||||||||||||||||||||||||||||||||
| name: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| created_at: datetime | ||||||||||||||||||||||||||||||||||||||||||||||||||
| updated_at: Optional[datetime] | ||||||||||||||||||||||||||||||||||||||||||||||||||
| extension: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| mime_type: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| raw_data_location: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| def get_datasets_router() -> APIRouter: | ||||||||||||||||||||||||||||||||||||||||||||||||||
| router = APIRouter() | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| class ErrorResponseDTO(BaseModel): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| message: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| class DatasetDTO(OutDTO): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| id: UUID | ||||||||||||||||||||||||||||||||||||||||||||||||||
| name: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| created_at: datetime | ||||||||||||||||||||||||||||||||||||||||||||||||||
| updated_at: Optional[datetime] | ||||||||||||||||||||||||||||||||||||||||||||||||||
| owner_id: UUID | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| @router.get("/", response_model=list[DatasetDTO]) | ||||||||||||||||||||||||||||||||||||||||||||||||||
| async def get_datasets(user: User = Depends(get_authenticated_user)): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| try: | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -96,15 +106,6 @@ async def get_dataset_graph(dataset_id: str, user: User = Depends(get_authentica | |||||||||||||||||||||||||||||||||||||||||||||||||
| content="Graphistry credentials are not set. Please set them in your .env file.", | ||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
+103
to
+107
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Replace bare except with specific exception handling. Using bare except clauses can mask important errors and make debugging difficult. return JSONResponse(
status_code=200,
content=str(graph_url),
)
- except:
+ except (ConnectionError, ValueError) as error:
+ logger.error(f"Error generating graph: {str(error)}")
return JSONResponse(
status_code=409,
content="Graphistry credentials are not set. Please set them in your .env file.",
)
🧰 Tools🪛 Ruff93-93: Do not use bare (E722)
Comment on lines
+95
to
+107
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Improve error handling in graph generation. The current implementation has two issues:
Apply this diff to improve error handling: try:
graph_client = await get_graph_engine()
graph_url = await render_graph(graph_client.graph)
- return JSONResponse(
- status_code=200,
- content=str(graph_url),
- )
+ return str(graph_url)
- except:
+ except (ConnectionError, ValueError) as error:
+ logger.error(f"Error generating graph: {str(error)}")
+ raise HTTPException(
+ status_code=409,
+ detail="Graphistry credentials are not set. Please set them in your .env file."
+ ) from error
- return JSONResponse(
- status_code=409,
- content="Graphistry credentials are not set. Please set them in your .env file.",
- )📝 Committable suggestion
Suggested change
🧰 Tools🪛 Ruff103-103: Do not use bare (E722) |
||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| class DataDTO(OutDTO): | ||||||||||||||||||||||||||||||||||||||||||||||||||
| id: UUID | ||||||||||||||||||||||||||||||||||||||||||||||||||
| name: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| created_at: datetime | ||||||||||||||||||||||||||||||||||||||||||||||||||
| updated_at: Optional[datetime] | ||||||||||||||||||||||||||||||||||||||||||||||||||
| extension: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| mime_type: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
| raw_data_location: str | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| @router.get("/{dataset_id}/data", response_model=list[DataDTO], | ||||||||||||||||||||||||||||||||||||||||||||||||||
| responses={404: {"model": ErrorResponseDTO}}) | ||||||||||||||||||||||||||||||||||||||||||||||||||
| async def get_dataset_data(dataset_id: str, user: User = Depends(get_authenticated_user)): | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
@@ -157,18 +158,21 @@ async def get_raw_data(dataset_id: str, data_id: str, user: User = Depends(get_a | |||||||||||||||||||||||||||||||||||||||||||||||||
| dataset_data = await get_dataset_data(dataset.id) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| if dataset_data is None: | ||||||||||||||||||||||||||||||||||||||||||||||||||
| raise HTTPException(status_code=404, detail=f"Dataset ({dataset_id}) not found.") | ||||||||||||||||||||||||||||||||||||||||||||||||||
| raise HTTPException(status_code=404, detail=f"No data found in dataset ({dataset_id}).") | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| data = [data for data in dataset_data if str(data.id) == data_id][0] | ||||||||||||||||||||||||||||||||||||||||||||||||||
| matching_data = [data for data in dataset_data if str(data.id) == data_id] | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| if data is None: | ||||||||||||||||||||||||||||||||||||||||||||||||||
| # Check if matching_data contains an element | ||||||||||||||||||||||||||||||||||||||||||||||||||
| if len(matching_data) == 0: | ||||||||||||||||||||||||||||||||||||||||||||||||||
| return JSONResponse( | ||||||||||||||||||||||||||||||||||||||||||||||||||
| status_code=404, | ||||||||||||||||||||||||||||||||||||||||||||||||||
| content={ | ||||||||||||||||||||||||||||||||||||||||||||||||||
| "detail": f"Data ({data_id}) not found in dataset ({dataset_id})." | ||||||||||||||||||||||||||||||||||||||||||||||||||
| } | ||||||||||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| data = matching_data[0] | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| return data.raw_data_location | ||||||||||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||||||||||
| return router | ||||||||||||||||||||||||||||||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,12 +5,13 @@ | |
| from cognee.api.DTO import InDTO | ||
| from cognee.modules.users.methods import get_authenticated_user | ||
|
|
||
| def get_search_router(): | ||
| router = APIRouter() | ||
|
|
||
| class SearchPayloadDTO(InDTO): | ||
| search_type: SearchType | ||
| query: str | ||
| class SearchPayloadDTO(InDTO): | ||
| search_type: SearchType | ||
| query: str | ||
|
|
||
| def get_search_router() -> APIRouter: | ||
| router = APIRouter() | ||
|
|
||
| @router.post("/", response_model = list) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion Specify the response model type. The from pydantic import BaseModel
class SearchResultDTO(BaseModel):
# Add appropriate fields based on the actual search results
id: str
title: str
# ... other fields
@router.post("/", response_model=list[SearchResultDTO]) |
||
| async def search(payload: SearchPayloadDTO, user: User = Depends(get_authenticated_user)): | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -4,41 +4,40 @@ | |||||||||||||||||||||||||
| from cognee.modules.users.methods import get_authenticated_user | ||||||||||||||||||||||||||
| from fastapi import Depends | ||||||||||||||||||||||||||
| from cognee.modules.users.models import User | ||||||||||||||||||||||||||
| from cognee.modules.settings.get_settings import LLMConfig, VectorDBConfig | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| def get_settings_router(): | ||||||||||||||||||||||||||
| router = APIRouter() | ||||||||||||||||||||||||||
| class LLMConfigOutputDTO(OutDTO, LLMConfig): | ||||||||||||||||||||||||||
| pass | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| class VectorDBConfigOutputDTO(OutDTO, VectorDBConfig): | ||||||||||||||||||||||||||
| pass | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| from cognee.modules.settings.get_settings import LLMConfig, VectorDBConfig | ||||||||||||||||||||||||||
| class SettingsDTO(OutDTO): | ||||||||||||||||||||||||||
| llm: LLMConfigOutputDTO | ||||||||||||||||||||||||||
| vector_db: VectorDBConfigOutputDTO | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| class LLMConfigDTO(OutDTO, LLMConfig): | ||||||||||||||||||||||||||
| pass | ||||||||||||||||||||||||||
| class LLMConfigInputDTO(InDTO): | ||||||||||||||||||||||||||
| provider: Union[Literal["openai"], Literal["ollama"], Literal["anthropic"]] | ||||||||||||||||||||||||||
| model: str | ||||||||||||||||||||||||||
| api_key: str | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| class VectorDBConfigDTO(OutDTO, VectorDBConfig): | ||||||||||||||||||||||||||
| pass | ||||||||||||||||||||||||||
| class VectorDBConfigInputDTO(InDTO): | ||||||||||||||||||||||||||
| provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]] | ||||||||||||||||||||||||||
| url: str | ||||||||||||||||||||||||||
| api_key: str | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| class SettingsDTO(OutDTO): | ||||||||||||||||||||||||||
| llm: LLMConfigDTO | ||||||||||||||||||||||||||
| vector_db: VectorDBConfigDTO | ||||||||||||||||||||||||||
| class SettingsPayloadDTO(InDTO): | ||||||||||||||||||||||||||
| llm: Optional[LLMConfigInputDTO] = None | ||||||||||||||||||||||||||
| vector_db: Optional[VectorDBConfigInputDTO] = None | ||||||||||||||||||||||||||
|
Comment on lines
+19
to
+31
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion Add field validation for api_key and url. Consider adding Pydantic field validators to ensure:
Example implementation: from pydantic import validator, HttpUrl
class LLMConfigInputDTO(InDTO):
provider: Literal["openai", "ollama", "anthropic"]
model: str
api_key: str
@validator('api_key')
def validate_api_key(cls, v):
if len(v) < 8:
raise ValueError('API key must be at least 8 characters')
return v
class VectorDBConfigInputDTO(InDTO):
provider: Literal["lancedb", "qdrant", "weaviate", "pgvector"]
url: HttpUrl # This will automatically validate URL format
api_key: str
@validator('api_key')
def validate_api_key(cls, v):
if len(v) < 8:
raise ValueError('API key must be at least 8 characters')
return v |
||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| def get_settings_router() -> APIRouter: | ||||||||||||||||||||||||||
| router = APIRouter() | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| @router.get("/", response_model=SettingsDTO) | ||||||||||||||||||||||||||
| async def get_settings(user: User = Depends(get_authenticated_user)): | ||||||||||||||||||||||||||
| from cognee.modules.settings import get_settings as get_cognee_settings | ||||||||||||||||||||||||||
| return get_cognee_settings() | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
|
Comment on lines
+36
to
+40
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Add error handling to GET endpoint. The endpoint should handle potential errors from Apply this diff: @router.get("/", response_model=SettingsDTO)
async def get_settings(user: User = Depends(get_authenticated_user)):
+ from fastapi import HTTPException
from cognee.modules.settings import get_settings as get_cognee_settings
- return get_cognee_settings()
+ try:
+ return get_cognee_settings()
+ except Exception as e:
+ raise HTTPException(status_code=500, detail=f"Failed to retrieve settings: {str(e)}")📝 Committable suggestion
Suggested change
|
||||||||||||||||||||||||||
| class LLMConfigDTO(InDTO): | ||||||||||||||||||||||||||
| provider: Union[Literal["openai"], Literal["ollama"], Literal["anthropic"]] | ||||||||||||||||||||||||||
| model: str | ||||||||||||||||||||||||||
| api_key: str | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| class VectorDBConfigDTO(InDTO): | ||||||||||||||||||||||||||
| provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]] | ||||||||||||||||||||||||||
| url: str | ||||||||||||||||||||||||||
| api_key: str | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| class SettingsPayloadDTO(InDTO): | ||||||||||||||||||||||||||
| llm: Optional[LLMConfigDTO] = None | ||||||||||||||||||||||||||
| vector_db: Optional[VectorDBConfigDTO] = None | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
| @router.post("/", response_model=None) | ||||||||||||||||||||||||||
| async def save_settings(new_settings: SettingsPayloadDTO, user: User = Depends(get_authenticated_user)): | ||||||||||||||||||||||||||
| from cognee.modules.settings import save_llm_config, save_vector_db_config | ||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type Checking Error: Inconsistent Handling of 'data'
The condition
if isinstance(data, str) and data.startswith("http"):checks ifdatais a string, butdatamight be a list ofUploadFile. This could cause unexpected behavior.Ensure that
datais properly handled for both string URLs and lists ofUploadFile. After updating the type annotation, adjust the logic to handle each case appropriately.