Coverage for src/semware/api/data.py: 84%

1"""API endpoints for data operations.""" 

2 

3from fastapi import APIRouter, Depends, HTTPException, status 

4from loguru import logger 

5 

6from ..models.requests import DeleteDataResponse, UpsertDataRequest, UpsertDataResponse 

7from ..services.embedding import embedding_service 

8from ..services.vectordb import vectordb 

9from .auth import api_key_auth 

10 

11router = APIRouter() 

12 

13 

14@router.post( 

15 "/{table_name}/data", 

16 response_model=UpsertDataResponse, 

17 status_code=status.HTTP_201_CREATED, 

18 summary="Insert or update data records", 

19 description="Insert new records or update existing ones in the specified table. Embeddings are automatically generated.", 

20) 

21async def upsert_data( 

22 table_name: str, request: UpsertDataRequest, api_key: str = Depends(api_key_auth) 

23) -> UpsertDataResponse: 

24 """Insert or update data records in a table.""" 

25 try: 

26 logger.info(f"Upserting {len(request.records)} records in table: {table_name}") 

27 

28 # Get table schema 

29 try: 

30 schema = vectordb.get_table_schema(table_name) 

31 except ValueError as e: 

32 raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) 

33 

34 # Validate all records first 

35 texts_for_embedding = [] 

36 for i, record in enumerate(request.records): 

37 try: 

38 # Validate record structure 

39 vectordb._validate_record(record, schema) 

40 

41 # Extract text for embedding 

42 text = record.get_text_for_embedding(schema.embedding_column) 

43 texts_for_embedding.append(text) 

44 

45 except ValueError as e: 

46 logger.error(f"Record {i} validation failed: {e}") 

47 raise HTTPException( 

48 status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, 

49 detail=f"Record {i}: {str(e)}", 

50 ) 

51 

52 # Generate embeddings 

53 logger.info(f"Generating embeddings for {len(texts_for_embedding)} records") 

54 embeddings = embedding_service.generate_embeddings(texts_for_embedding) 

55 

56 # Upsert records with embeddings 

57 inserted_count, updated_count = vectordb.upsert_records( 

58 table_name, request.records, embeddings 

59 ) 

60 

61 return UpsertDataResponse( 

62 message=f"Successfully processed {len(request.records)} records", 

63 inserted_count=inserted_count, 

64 updated_count=updated_count, 

65 ) 

66 

67 except HTTPException: 

68 raise 

69 except Exception as e: 

70 logger.exception(f"Error upserting data in table '{table_name}': {e}") 

71 raise HTTPException( 

72 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

73 detail="Failed to upsert data", 

74 ) 

75 

76 

77@router.delete( 

78 "/{table_name}/data/{record_id}", 

79 response_model=DeleteDataResponse, 

80 summary="Delete a data record", 

81 description="Delete a specific record from the table by its ID.", 

82) 

83async def delete_data( 

84 table_name: str, record_id: str, api_key: str = Depends(api_key_auth) 

85) -> DeleteDataResponse: 

86 """Delete a data record from a table.""" 

87 try: 

88 logger.info(f"Deleting record '{record_id}' from table: {table_name}") 

89 

90 # Check if table exists 

91 try: 

92 vectordb.get_table_schema(table_name) 

93 except ValueError as e: 

94 raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) 

95 

96 # Check if record exists 

97 existing_record = vectordb.get_record(table_name, record_id) 

98 if existing_record is None: 

99 raise HTTPException( 

100 status_code=status.HTTP_404_NOT_FOUND, 

101 detail=f"Record with ID '{record_id}' not found in table '{table_name}'", 

102 ) 

103 

104 # Delete the record 

105 vectordb.delete_record(table_name, record_id) 

106 

107 return DeleteDataResponse( 

108 message="Record deleted successfully", deleted_id=record_id 

109 ) 

110 

111 except HTTPException: 

112 raise 

113 except Exception as e: 

114 logger.exception( 

115 f"Error deleting record '{record_id}' from table '{table_name}': {e}" 

116 ) 

117 raise HTTPException( 

118 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

119 detail="Failed to delete record", 

120 ) 

121 

122 

123@router.get( 

124 "/{table_name}/data/{record_id}", 

125 summary="Get a data record", 

126 description="Retrieve a specific record from the table by its ID.", 

127) 

128async def get_data( 

129 table_name: str, record_id: str, api_key: str = Depends(api_key_auth) 

130) -> dict: 

131 """Get a data record from a table.""" 

132 try: 

133 logger.debug(f"Getting record '{record_id}' from table: {table_name}") 

134 

135 # Check if table exists 

136 try: 

137 vectordb.get_table_schema(table_name) 

138 except ValueError as e: 

139 raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) 

140 

141 # Get the record 

142 record = vectordb.get_record(table_name, record_id) 

143 

144 if record is None: 

145 raise HTTPException( 

146 status_code=status.HTTP_404_NOT_FOUND, 

147 detail=f"Record with ID '{record_id}' not found in table '{table_name}'", 

148 ) 

149 

150 return {"table_name": table_name, "record_id": record_id, "data": record} 

151 

152 except HTTPException: 

153 raise 

154 except Exception as e: 

155 logger.exception( 

156 f"Error getting record '{record_id}' from table '{table_name}': {e}" 

157 ) 

158 raise HTTPException( 

159 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, 

160 detail="Failed to get record", 

161 )