diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..d3b7366 --- /dev/null +++ b/.env.example @@ -0,0 +1,21 @@ +# ============================================================ +# Decompile-AI — Environment Variables +# ============================================================ +# Copy this file to .env and fill in your values. +# .env is git-ignored and never committed. + +# --- LLM Provider --- + +# Default chat provider (ollama) +AI_DEFAULT_CHAT_PROVIDER=ollama + +# --- Model overrides (optional) --- + +# Ollama chat model (default: gemma4:12b) +# AI_CHAT_MODEL_OLLAMA=gemma4:12b + +# Ollama embedding model (default: qwen3-embedding:latest) +# AI_EMBEDDING_MODEL=qwen3-embedding:latest + +# Embedding vector dimension (default: 4096) +# AI_EMBEDDING_DIMENSION=4096 diff --git a/pom.xml b/pom.xml index 98fc4ef..83a30ef 100644 --- a/pom.xml +++ b/pom.xml @@ -26,10 +26,11 @@ - - 25 - 2.0.6 - + + 25 + 2.0.6 + 1.0.1 + org.springframework.boot @@ -67,7 +68,6 @@ org.postgresql postgresql - runtime org.springframework.modulith @@ -135,25 +135,42 @@ docker-java-transport-httpclient5 3.7.1 - - org.apache.commons - commons-lang3 - 3.20.0 - compile - - + + org.apache.commons + commons-lang3 + 3.20.0 + compile + + + org.springframework.ai + spring-ai-ollama + + + io.swagger.core.v3 + swagger-annotations-jakarta + + + + - - - - org.springframework.modulith - spring-modulith-bom - ${spring-modulith.version} - pom - import - - - + + + + org.springframework.modulith + spring-modulith-bom + ${spring-modulith.version} + pom + import + + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + + + diff --git a/src/main/java/ai/decompile/ai/config/AiConfig.java b/src/main/java/ai/decompile/ai/config/AiConfig.java new file mode 100644 index 0000000..b34e28d --- /dev/null +++ b/src/main/java/ai/decompile/ai/config/AiConfig.java @@ -0,0 +1,61 @@ +package ai.decompile.ai.config; + +import io.micrometer.observation.ObservationRegistry; +import org.springframework.ai.model.tool.ToolCallingManager; +import org.springframework.ai.ollama.OllamaChatModel; +import org.springframework.ai.ollama.OllamaEmbeddingModel; +import org.springframework.ai.ollama.api.OllamaApi; +import org.springframework.ai.ollama.api.OllamaOptions; +import org.springframework.ai.ollama.management.ModelManagementOptions; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.resilience.annotation.EnableResilientMethods; + +@Configuration +@EnableResilientMethods +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class AiConfig { + @Value("${spring.ai.ollama.base-url:http://localhost:11434}") + private String ollamaBaseUrl; + + @Value("${spring.ai.ollama.chat.options.model:gemma4:12b}") + private String ollamaChatModelName; + + @Value("${spring.ai.ollama.chat.options.temperature:0.3}") + private double ollamaTemperature; + + @Value("${spring.ai.ollama.embedding.options.model:qwen3-embedding:latest}") + private String ollamaEmbeddingModel; + + @Bean + OllamaApi ollamaApi() { + return OllamaApi.builder().baseUrl(ollamaBaseUrl).build(); + } + + @Bean + OllamaChatModel ollamaChatModel(OllamaApi ollamaApi, ObservationRegistry observationRegistry) { + return OllamaChatModel.builder() + .ollamaApi(ollamaApi) + .defaultOptions( + OllamaOptions.builder() + .model(ollamaChatModelName) + .temperature(ollamaTemperature) + .build()) + .toolCallingManager(ToolCallingManager.builder().build()) + .observationRegistry(observationRegistry) + .build(); + } + + @Bean + OllamaEmbeddingModel ollamaEmbeddingModel( + OllamaApi ollamaApi, ObservationRegistry observationRegistry) { + return OllamaEmbeddingModel.builder() + .ollamaApi(ollamaApi) + .defaultOptions(OllamaOptions.builder().model(ollamaEmbeddingModel).build()) + .observationRegistry(observationRegistry) + .modelManagementOptions(ModelManagementOptions.defaults()) + .build(); + } +} diff --git a/src/main/java/ai/decompile/ai/controller/ChatController.java b/src/main/java/ai/decompile/ai/controller/ChatController.java new file mode 100644 index 0000000..32b3b45 --- /dev/null +++ b/src/main/java/ai/decompile/ai/controller/ChatController.java @@ -0,0 +1,115 @@ +package ai.decompile.ai.controller; + +import ai.decompile.ai.model.dto.ChatRequest; +import ai.decompile.ai.model.dto.MessageResponse; +import ai.decompile.ai.model.dto.SessionResponse; +import ai.decompile.ai.service.ChatService; +import ai.decompile.ai.service.ChatSessionService; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.swagger.v3.oas.annotations.Hidden; +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.tags.Tag; +import jakarta.validation.Valid; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import lombok.RequiredArgsConstructor; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.http.codec.ServerSentEvent; +import org.springframework.web.bind.annotation.DeleteMapping; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.PathVariable; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestParam; +import org.springframework.web.bind.annotation.RestController; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; + +@RestController +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +@Tag(name = "Chat", description = "AI-powered chat for binary analysis") +public class ChatController { + private final ChatService chatService; + private final ChatSessionService sessionService; + private final ObjectMapper objectMapper; + + @PostMapping("/binaries/{binaryId}/chat/sessions") + @Operation(summary = "Create a new chat session for a binary") + public ResponseEntity createSession( + @PathVariable UUID binaryId, + @RequestParam(required = false) UUID projectId, + @RequestParam(required = false, defaultValue = "gemma4:12b") String model) { + if (projectId == null) { + return ResponseEntity.badRequest().build(); + } + var session = sessionService.createSession(binaryId, projectId, model); + return ResponseEntity.status(HttpStatus.CREATED).body(SessionResponse.from(session)); + } + + @GetMapping("/binaries/{binaryId}/chat/sessions") + @Operation(summary = "List chat sessions for a binary") + public ResponseEntity> listSessions(@PathVariable UUID binaryId) { + var sessions = sessionService.getSessions(binaryId); + return ResponseEntity.ok(sessions.stream().map(SessionResponse::from).toList()); + } + + @GetMapping("/chat/sessions/{sessionId}") + @Operation(summary = "Get chat session details") + public ResponseEntity getSession(@PathVariable UUID sessionId) { + var session = sessionService.getSession(sessionId); + return ResponseEntity.ok(SessionResponse.from(session)); + } + + @DeleteMapping("/chat/sessions/{sessionId}") + @Operation(summary = "Delete a chat session and all its messages") + public ResponseEntity deleteSession(@PathVariable UUID sessionId) { + sessionService.deleteSession(sessionId); + return ResponseEntity.noContent().build(); + } + + @GetMapping("/chat/sessions/{sessionId}/messages") + @Operation(summary = "Get message history for a chat session") + public ResponseEntity> getMessages(@PathVariable UUID sessionId) { + var messages = sessionService.getMessages(sessionId); + return ResponseEntity.ok(messages.stream().map(MessageResponse::from).toList()); + } + + @Hidden + @PostMapping( + value = "/chat/sessions/{sessionId}/messages", + produces = MediaType.TEXT_EVENT_STREAM_VALUE) + public Flux> sendMessage( + @PathVariable UUID sessionId, @Valid @RequestBody ChatRequest request) { + + return chatService + .chat(sessionId, request.content()) + .map(token -> event("chunk", "content", token)) + .concatWith(Mono.just(event("done", "sessionId", sessionId.toString()))) + .onErrorResume(e -> Flux.just(errorEvent(e))); + } + + private ServerSentEvent event(String type, String key, Object value) { + return ServerSentEvent.builder().data(toJson(Map.of("type", type, key, value))).build(); + } + + private ServerSentEvent errorEvent(Throwable e) { + String msg = e.getMessage() != null ? e.getMessage() : "Unknown error"; + return ServerSentEvent.builder() + .data(toJson(Map.of("type", "error", "message", msg))) + .build(); + } + + private String toJson(Object obj) { + try { + return objectMapper.writeValueAsString(obj); + } catch (JsonProcessingException e) { + return "{\"type\":\"error\",\"message\":\"JSON serialization failed\"}"; + } + } +} diff --git a/src/main/java/ai/decompile/ai/controller/package-info.java b/src/main/java/ai/decompile/ai/controller/package-info.java new file mode 100644 index 0000000..bc0ea4c --- /dev/null +++ b/src/main/java/ai/decompile/ai/controller/package-info.java @@ -0,0 +1 @@ +package ai.decompile.ai.controller; diff --git a/src/main/java/ai/decompile/ai/model/dto/ChatRequest.java b/src/main/java/ai/decompile/ai/model/dto/ChatRequest.java new file mode 100644 index 0000000..2e67bd9 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/dto/ChatRequest.java @@ -0,0 +1,6 @@ +package ai.decompile.ai.model.dto; + +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.Size; + +public record ChatRequest(@NotBlank @Size(min = 1, max = 4000) String content) {} diff --git a/src/main/java/ai/decompile/ai/model/dto/MessageResponse.java b/src/main/java/ai/decompile/ai/model/dto/MessageResponse.java new file mode 100644 index 0000000..8da29d1 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/dto/MessageResponse.java @@ -0,0 +1,18 @@ +package ai.decompile.ai.model.dto; + +import ai.decompile.ai.model.entity.ChatMessage; +import java.time.Instant; +import java.util.UUID; + +public record MessageResponse( + UUID id, UUID sessionId, String role, String content, Integer tokenCount, Instant createdAt) { + public static MessageResponse from(ChatMessage message) { + return new MessageResponse( + message.getId(), + message.getSession().getId(), + message.getRole(), + message.getContent(), + message.getTokenCount(), + message.getCreatedAt()); + } +} diff --git a/src/main/java/ai/decompile/ai/model/dto/SessionResponse.java b/src/main/java/ai/decompile/ai/model/dto/SessionResponse.java new file mode 100644 index 0000000..581ac40 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/dto/SessionResponse.java @@ -0,0 +1,27 @@ +package ai.decompile.ai.model.dto; + +import ai.decompile.ai.model.entity.ChatSession; +import java.time.Instant; +import java.util.UUID; + +public record SessionResponse( + UUID id, + UUID binaryId, + UUID projectId, + String title, + String chatModel, + int messageCount, + Instant createdAt, + Instant updatedAt) { + public static SessionResponse from(ChatSession session) { + return new SessionResponse( + session.getId(), + session.getBinaryId(), + session.getProjectId(), + session.getTitle(), + session.getChatModel(), + session.getMessages() != null ? session.getMessages().size() : 0, + session.getCreatedAt(), + session.getUpdatedAt()); + } +} diff --git a/src/main/java/ai/decompile/ai/model/dto/package-info.java b/src/main/java/ai/decompile/ai/model/dto/package-info.java new file mode 100644 index 0000000..0506e4c --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/dto/package-info.java @@ -0,0 +1 @@ +package ai.decompile.ai.model.dto; diff --git a/src/main/java/ai/decompile/ai/model/entity/ChatMessage.java b/src/main/java/ai/decompile/ai/model/entity/ChatMessage.java new file mode 100644 index 0000000..7313067 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/entity/ChatMessage.java @@ -0,0 +1,37 @@ +package ai.decompile.ai.model.entity; + +import jakarta.persistence.*; +import java.time.Instant; +import java.util.UUID; +import lombok.*; +import org.hibernate.annotations.CreationTimestamp; + +@Entity +@Table(name = "chat_message") +@Getter +@Setter +@NoArgsConstructor +@AllArgsConstructor +@Builder +public class ChatMessage { + @Id + @GeneratedValue(strategy = GenerationType.UUID) + private UUID id; + + @ManyToOne(fetch = FetchType.LAZY) + @JoinColumn(name = "session_id", nullable = false) + private ChatSession session; + + @Column(nullable = false, length = 20) + private String role; + + @Column(nullable = false, columnDefinition = "TEXT") + private String content; + + @Column(name = "token_count") + private Integer tokenCount; + + @CreationTimestamp + @Column(name = "created_at", nullable = false, updatable = false) + private Instant createdAt; +} diff --git a/src/main/java/ai/decompile/ai/model/entity/ChatSession.java b/src/main/java/ai/decompile/ai/model/entity/ChatSession.java new file mode 100644 index 0000000..79c9bec --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/entity/ChatSession.java @@ -0,0 +1,60 @@ +package ai.decompile.ai.model.entity; + +import ai.decompile.workspace.model.entity.Binary; +import ai.decompile.workspace.model.entity.Project; +import jakarta.persistence.*; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; +import lombok.*; +import org.hibernate.annotations.CreationTimestamp; +import org.hibernate.annotations.UpdateTimestamp; + +@Entity +@Table(name = "chat_session") +@Getter +@Setter +@NoArgsConstructor +@AllArgsConstructor +@Builder +public class ChatSession { + @Id + @GeneratedValue(strategy = GenerationType.UUID) + private UUID id; + + @Column(name = "binary_id", nullable = false) + private UUID binaryId; + + @ManyToOne(fetch = FetchType.LAZY) + @JoinColumn(name = "binary_id", insertable = false, updatable = false) + private Binary binary; + + @Column(name = "project_id", nullable = false) + private UUID projectId; + + @ManyToOne(fetch = FetchType.LAZY) + @JoinColumn(name = "project_id", insertable = false, updatable = false) + private Project project; + + @Column private String title; + + @Column(name = "chat_model", nullable = false, length = 100) + private String chatModel; + + @OneToMany( + mappedBy = "session", + cascade = CascadeType.ALL, + orphanRemoval = true, + fetch = FetchType.LAZY) + @Builder.Default + private List messages = new ArrayList<>(); + + @CreationTimestamp + @Column(name = "created_at", nullable = false, updatable = false) + private Instant createdAt; + + @UpdateTimestamp + @Column(name = "updated_at", nullable = false) + private Instant updatedAt; +} diff --git a/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java b/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java new file mode 100644 index 0000000..dcaebed --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java @@ -0,0 +1,55 @@ +package ai.decompile.ai.model.entity; + +import ai.decompile.ai.model.jdbc.PostgreSQLVectorJdbcType; +import jakarta.persistence.*; +import java.time.Instant; +import java.util.UUID; +import lombok.*; +import org.hibernate.annotations.CreationTimestamp; +import org.hibernate.annotations.JdbcTypeCode; +import org.hibernate.annotations.JdbcTypeRegistration; +import org.hibernate.type.SqlTypes; + +@Entity +@Table(name = "embedding_chunk") +@Getter +@Setter +@NoArgsConstructor +@AllArgsConstructor +@Builder +@JdbcTypeRegistration(PostgreSQLVectorJdbcType.class) +public class EmbeddingChunk { + @Id + @GeneratedValue(strategy = GenerationType.UUID) + private UUID id; + + @Column(name = "binary_id", nullable = false) + private UUID binaryId; + + @Column(name = "chunk_type", nullable = false, length = 30) + private String chunkType; + + @Column(name = "source_type", length = 30) + private String sourceType; + + @Column(name = "source_id") + private UUID sourceId; + + @Column(nullable = false, columnDefinition = "TEXT") + private String content; + + @JdbcTypeCode(SqlTypes.VECTOR) + @Column(name = "embedding", columnDefinition = "vector(4096)") + private float[] embedding; + + @JdbcTypeCode(SqlTypes.JSON) + @Column(name = "metadata", columnDefinition = "jsonb") + private String metadata; + + @Column(name = "token_count") + private Integer tokenCount; + + @CreationTimestamp + @Column(name = "created_at", nullable = false, updatable = false) + private Instant createdAt; +} diff --git a/src/main/java/ai/decompile/ai/model/entity/package-info.java b/src/main/java/ai/decompile/ai/model/entity/package-info.java new file mode 100644 index 0000000..1b0bbf6 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/entity/package-info.java @@ -0,0 +1 @@ +package ai.decompile.ai.model.entity; diff --git a/src/main/java/ai/decompile/ai/model/jdbc/PostgreSQLVectorJdbcType.java b/src/main/java/ai/decompile/ai/model/jdbc/PostgreSQLVectorJdbcType.java new file mode 100644 index 0000000..c37c46d --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/jdbc/PostgreSQLVectorJdbcType.java @@ -0,0 +1,115 @@ +package ai.decompile.ai.model.jdbc; + +import java.sql.CallableStatement; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Types; +import java.util.Objects; +import org.hibernate.type.SqlTypes; +import org.hibernate.type.descriptor.ValueBinder; +import org.hibernate.type.descriptor.ValueExtractor; +import org.hibernate.type.descriptor.WrapperOptions; +import org.hibernate.type.descriptor.java.JavaType; +import org.hibernate.type.descriptor.jdbc.BasicBinder; +import org.hibernate.type.descriptor.jdbc.BasicExtractor; +import org.hibernate.type.descriptor.jdbc.JdbcType; +import org.postgresql.util.PGobject; + +public class PostgreSQLVectorJdbcType implements JdbcType { + @Override + public int getJdbcTypeCode() { + return Types.OTHER; + } + + @Override + public int getDefaultSqlTypeCode() { + return SqlTypes.VECTOR; + } + + @Override + public String toString() { + return "PostgreSQLVectorJdbcType"; + } + + @Override + public ValueBinder getBinder(JavaType javaType) { + return new BasicBinder<>(javaType, this) { + @Override + protected void doBind(PreparedStatement st, X value, int index, WrapperOptions options) + throws SQLException { + var pgObj = new PGobject(); + pgObj.setType("vector"); + pgObj.setValue(formatVector((float[]) value)); + st.setObject(index, pgObj); + } + + @Override + protected void doBind(CallableStatement st, X value, String name, WrapperOptions options) + throws SQLException { + var pgObj = new PGobject(); + pgObj.setType("vector"); + pgObj.setValue(formatVector((float[]) value)); + st.setObject(name, pgObj); + } + }; + } + + @Override + public ValueExtractor getExtractor(JavaType javaType) { + return new BasicExtractor<>(javaType, this) { + @Override + protected X doExtract(ResultSet rs, int paramIndex, WrapperOptions options) + throws SQLException { + return javaType.wrap(parseVector(rs.getString(paramIndex)), options); + } + + @Override + protected X doExtract(CallableStatement statement, int index, WrapperOptions options) + throws SQLException { + return javaType.wrap(parseVector(statement.getString(index)), options); + } + + @Override + protected X doExtract(CallableStatement statement, String name, WrapperOptions options) + throws SQLException { + return javaType.wrap(parseVector(statement.getString(name)), options); + } + }; + } + + static float[] parseVector(String vectorStr) { + if (Objects.isNull(vectorStr) || vectorStr.isBlank()) { + return new float[0]; + } + + var trimmed = vectorStr.trim(); + + if (trimmed.startsWith("[") && trimmed.endsWith("]")) { + trimmed = trimmed.substring(1, trimmed.length() - 1); + } + + var parts = trimmed.split(","); + var result = new float[parts.length]; + for (int i = 0; i < parts.length; i++) { + result[i] = Float.parseFloat(parts[i].trim()); + } + return result; + } + + static String formatVector(float[] vector) { + if (Objects.isNull(vector) || vector.length == 0) { + return "[]"; + } + + var sb = new StringBuilder("["); + for (int i = 0; i < vector.length; i++) { + if (i > 0) { + sb.append(","); + } + + sb.append(vector[i]); + } + return sb.append("]").toString(); + } +} diff --git a/src/main/java/ai/decompile/ai/model/repository/ChatMessageRepository.java b/src/main/java/ai/decompile/ai/model/repository/ChatMessageRepository.java new file mode 100644 index 0000000..4df2413 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/repository/ChatMessageRepository.java @@ -0,0 +1,10 @@ +package ai.decompile.ai.model.repository; + +import ai.decompile.ai.model.entity.ChatMessage; +import java.util.List; +import java.util.UUID; +import org.springframework.data.jpa.repository.JpaRepository; + +public interface ChatMessageRepository extends JpaRepository { + List findBySessionIdOrderByCreatedAtAsc(UUID sessionId); +} diff --git a/src/main/java/ai/decompile/ai/model/repository/ChatSessionRepository.java b/src/main/java/ai/decompile/ai/model/repository/ChatSessionRepository.java new file mode 100644 index 0000000..fa2fd09 --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/repository/ChatSessionRepository.java @@ -0,0 +1,10 @@ +package ai.decompile.ai.model.repository; + +import ai.decompile.ai.model.entity.ChatSession; +import java.util.List; +import java.util.UUID; +import org.springframework.data.jpa.repository.JpaRepository; + +public interface ChatSessionRepository extends JpaRepository { + List findByBinaryIdOrderByUpdatedAtDesc(UUID binaryId); +} diff --git a/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java b/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java new file mode 100644 index 0000000..54501dd --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java @@ -0,0 +1,65 @@ +package ai.decompile.ai.model.repository; + +import ai.decompile.ai.model.entity.EmbeddingChunk; +import java.util.List; +import java.util.UUID; +import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.data.jpa.repository.Modifying; +import org.springframework.data.jpa.repository.Query; +import org.springframework.data.repository.query.Param; +import org.springframework.transaction.annotation.Transactional; + +public interface EmbeddingChunkRepository extends JpaRepository { + @Query( + value = + """ + SELECT * FROM embedding_chunk + WHERE binary_id = :binaryId + AND 1 - (embedding <=> CAST(:queryEmbedding AS vector)) > :threshold + ORDER BY embedding <=> CAST(:queryEmbedding AS vector) + LIMIT :limit + """, + nativeQuery = true) + List findSimilar( + @Param("binaryId") UUID binaryId, + @Param("queryEmbedding") String queryEmbedding, + @Param("threshold") double threshold, + @Param("limit") int limit); + + @Modifying + @Transactional + @Query( + value = + """ + INSERT INTO embedding_chunk + (id, binary_id, chunk_type, source_type, source_id, content, embedding, metadata, token_count, created_at) + VALUES + (:id, :binaryId, :chunkType, :sourceType, :sourceId, :content, CAST(:embedding AS vector), CAST(:metadata AS jsonb), :tokenCount, :createdAt) + """, + nativeQuery = true) + void insertChunk( + @Param("id") UUID id, + @Param("binaryId") UUID binaryId, + @Param("chunkType") String chunkType, + @Param("sourceType") String sourceType, + @Param("sourceId") UUID sourceId, + @Param("content") String content, + @Param("embedding") String embedding, + @Param("metadata") String metadata, + @Param("tokenCount") Integer tokenCount, + @Param("createdAt") java.time.Instant createdAt); + + void deleteByBinaryId(UUID binaryId); + + @Query( + value = + """ + SELECT * FROM embedding_chunk + WHERE binary_id = :binaryId + AND LOWER(content) LIKE LOWER(CONCAT('%', :term, '%')) + LIMIT :limit + """, + nativeQuery = true) + List searchByContent( + @Param("binaryId") UUID binaryId, @Param("term") String term, @Param("limit") int limit); +} diff --git a/src/main/java/ai/decompile/ai/model/repository/package-info.java b/src/main/java/ai/decompile/ai/model/repository/package-info.java new file mode 100644 index 0000000..adffdcb --- /dev/null +++ b/src/main/java/ai/decompile/ai/model/repository/package-info.java @@ -0,0 +1 @@ +package ai.decompile.ai.model.repository; diff --git a/src/main/java/ai/decompile/ai/package-info.java b/src/main/java/ai/decompile/ai/package-info.java new file mode 100644 index 0000000..3e75553 --- /dev/null +++ b/src/main/java/ai/decompile/ai/package-info.java @@ -0,0 +1,10 @@ +@org.springframework.modulith.ApplicationModule( + displayName = "AI", + allowedDependencies = { + "workspace::entities", + "workspace::services", + "analysis::entities", + "analysis::services", + "common" + }) +package ai.decompile.ai; diff --git a/src/main/java/ai/decompile/ai/service/AiUtils.java b/src/main/java/ai/decompile/ai/service/AiUtils.java new file mode 100644 index 0000000..85b4bec --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/AiUtils.java @@ -0,0 +1,50 @@ +package ai.decompile.ai.service; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.Objects; +import lombok.AccessLevel; +import lombok.NoArgsConstructor; +import lombok.extern.log4j.Log4j2; + +@Log4j2 +@NoArgsConstructor(access = AccessLevel.PRIVATE) +public final class AiUtils { + public static String truncate(String text, int maxChars) { + if (Objects.isNull(text) || text.length() <= maxChars) { + return Objects.nonNull(text) ? text : ""; + } + + return text.substring(0, maxChars) + "\n... [truncated]"; + } + + public static String orUnknown(String value) { + return Objects.nonNull(value) ? value : "unknown"; + } + + public static int estimateTokens(String text) { + return Objects.nonNull(text) ? (int) Math.ceil(text.length() / 3.5) : 0; + } + + public static String toVectorString(float[] embedding) { + var sb = new StringBuilder("["); + for (var i = 0; i < embedding.length; i++) { + if (i > 0) { + sb.append(","); + } + + sb.append(embedding[i]); + } + + return sb.append("]").toString(); + } + + public static String toJson(ObjectMapper mapper, Object obj) { + try { + return mapper.writeValueAsString(obj); + } catch (JsonProcessingException e) { + log.warn("Failed to serialize object to JSON: {}", e.toString()); + return "{}"; + } + } +} diff --git a/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java b/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java new file mode 100644 index 0000000..7b64c89 --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java @@ -0,0 +1,309 @@ +package ai.decompile.ai.service; + +import ai.decompile.ai.model.entity.EmbeddingChunk; +import ai.decompile.ai.model.repository.EmbeddingChunkRepository; +import ai.decompile.analysis.model.entity.StaticAnalysis; +import ai.decompile.analysis.model.entity.StaticFunction; +import ai.decompile.analysis.model.entity.StaticXref; +import ai.decompile.analysis.model.repository.StaticAnalysisRepository; +import ai.decompile.analysis.model.repository.StaticFunctionRepository; +import ai.decompile.analysis.model.repository.StaticXrefRepository; +import ai.decompile.workspace.model.entity.Binary; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Objects; +import java.util.Set; +import java.util.UUID; +import java.util.regex.Pattern; +import java.util.stream.Stream; +import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +@Log4j2 +@Component +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class ChatContextBuilder { + // Detects: IDA auto-names (sub_XXXX), hex addresses (0xXXXXXXXX), + // camelCase/SCREAMING_SNAKE names with 2+ uppercase segments, and offset labels (name@addr) + private static final Pattern FUNCTION_PATTERN = + Pattern.compile("\\b(sub_\\w+|0x[0-9a-fA-F]+|[A-Z]\\w*[A-Z]\\w*|\\w+@\\w+)\\b"); + private static final int MIN_TOKEN_LENGTH = 3; + private static final int MAX_FUNCTION_REFERENCES = 5; + private static final int NEIGHBOR_CONTEXT_MAX_CHARS = 3000; + private static final int DEFAULT_CALL_GRAPH_HOPS = 1; + + private final EmbeddingService embeddingService; + private final EmbeddingChunkRepository embeddingChunkRepository; + private final StaticAnalysisRepository analysisRepository; + private final StaticFunctionRepository functionRepository; + private final StaticXrefRepository xrefRepository; + + @Value("${app.ai.rag.top-k:10}") + private int topK; + + @Value("${app.ai.rag.similarity-threshold:0.6}") + private double similarityThreshold; + + @Value("${app.ai.rag.fallback-min-ratio:0.4}") + private double fallbackMinRatio; + + @Value("${app.ai.rag.function-context-max-chars:5000}") + private int functionContextMaxChars; + + public String buildContext(UUID binaryId, String query, Binary binary) { + var merged = new LinkedHashSet(); + var inlineSections = new ArrayList(); + + enrichWithSemanticAndTextualSearch(merged, binaryId, query); + enrichWithCallGraph(inlineSections, merged, query, binaryId); + + return formatContext(inlineSections, merged, binary.getFilename()); + } + + public String buildSystemPrompt(String context, Binary binary) { + return String.format( + """ + You are an expert reverse engineer and malware analyst. + You are analyzing a binary file with the following characteristics. + Use the provided function disassembly and context to answer the user's questions. + When referencing code, mention function names, addresses, and the binary context. + If the provided context does not contain enough information, acknowledge the limitation + and suggest what additional analysis would help. + + [BINARY OVERVIEW] + Filename: %s + Format: %s + Architecture: %s + Compiler: %s + + %s + + [CONVERSATION GUIDELINES] + - Cite function addresses when relevant (e.g., "the function at 0x401000") + - Explain assembly in natural language when useful + - If you identify patterns (loops, conditionals, API calls), highlight them + - If you detect malicious intent, point it out objectively + - Be concise but thorough in your technical analysis + - If the context includes a [REQUESTED FUNCTION ANALYSIS] section with full assembly, + analyze that function in depth rather than saying you don't have access to it + """, + AiUtils.orUnknown(binary.getFilename()), + AiUtils.orUnknown(binary.getFormat()), + AiUtils.orUnknown(binary.getArchitecture()), + AiUtils.orUnknown(binary.getCompiler()), + context); + } + + private void enrichWithSemanticAndTextualSearch( + LinkedHashSet merged, UUID binaryId, String query) { + try { + var semanticChunks = semanticSearch(binaryId, query); + log.debug( + "Semantic search returned {} chunks for binary {}", semanticChunks.size(), binaryId); + merged.addAll(semanticChunks); + + if (semanticChunks.size() < topK * fallbackMinRatio) { + var tokens = extractMeaningfulTokens(query); + enrichWithTextualFallback(merged, binaryId, tokens); + } + } catch (Exception e) { + log.warn("Semantic/textual search failed, continuing with call graph: {}", e.getMessage()); + } + } + + private void enrichWithTextualFallback( + LinkedHashSet merged, UUID binaryId, List tokens) { + int remaining = topK - merged.size(); + + for (var token : tokens) { + if (remaining <= 0 || merged.size() >= topK) { + break; + } + + remaining = addTextChunksToMerged(merged, binaryId, token, remaining); + } + } + + private int addTextChunksToMerged( + LinkedHashSet merged, UUID binaryId, String token, int limit) { + var textChunks = embeddingChunkRepository.searchByContent(binaryId, token, limit); + + for (var chunk : textChunks) { + if (limit <= 0) { + break; + } + + if (merged.add(chunk)) { + limit--; + } + } + + return limit; + } + + private void enrichWithCallGraph( + List inlineSections, + LinkedHashSet merged, + String query, + UUID binaryId) { + try { + var detectedFunctions = detectFunctionReferences(query, binaryId); + + for (var func : detectedFunctions.limit(5).toList()) { + if (isFunctionAlreadyInContext(merged, func)) { + continue; + } + + var section = buildFunctionContext(func); + inlineSections.add(section); + } + } catch (Exception e) { + log.warn("Call graph expansion failed: {}", e.getMessage()); + } + } + + private boolean isFunctionAlreadyInContext( + LinkedHashSet merged, StaticFunction func) { + return merged.stream() + .anyMatch(c -> Objects.nonNull(c.getSourceId()) && c.getSourceId().equals(func.getId())); + } + + private String formatContext( + List inlineSections, LinkedHashSet merged, String filename) { + if (merged.isEmpty() && inlineSections.isEmpty()) { + return "[No relevant code sections found for this query in binary " + filename + ".]"; + } + + var sb = new StringBuilder(); + + if (!inlineSections.isEmpty()) { + sb.append("[REQUESTED FUNCTION ANALYSIS]\n\n"); + for (var section : inlineSections) { + sb.append(section).append("\n"); + } + } + + if (!merged.isEmpty()) { + sb.append("[RELEVANT CODE SECTIONS]\n\n"); + int idx = 1; + for (var chunk : merged) { + sb.append("Section ").append(idx++).append(":\n"); + sb.append(chunk.getContent()).append("\n"); + } + } + + return sb.toString(); + } + + private List semanticSearch(UUID binaryId, String query) { + var queryEmbedding = embeddingService.embed(query); + var vectorStr = AiUtils.toVectorString(queryEmbedding); + return embeddingChunkRepository.findSimilar(binaryId, vectorStr, similarityThreshold, topK); + } + + private List extractMeaningfulTokens(String query) { + return Stream.of(query.split("[\\s,.;:!?()\\[\\]{}\"']+")) + .filter(t -> t.length() >= MIN_TOKEN_LENGTH) + .distinct() + .limit(10) + .toList(); + } + + private Stream detectFunctionReferences(String query, UUID binaryId) { + var analysis = getLatestAnalysis(binaryId); + if (Objects.isNull(analysis)) { + return Stream.empty(); + } + + var tokens = extractMeaningfulTokens(query); + var matcher = FUNCTION_PATTERN.matcher(query); + var extraTokens = new ArrayList(); + while (matcher.find() && extraTokens.size() < 10) { + extraTokens.add(matcher.group(1)); + } + + var allTokens = Stream.concat(tokens.stream(), extraTokens.stream()).distinct().toList(); + + Set results = new LinkedHashSet<>(); + for (var token : allTokens) { + if (results.size() >= MAX_FUNCTION_REFERENCES) { + break; + } + + var functions = + functionRepository.findByNameIgnoreCaseContainingAndAnalysisId(token, analysis.getId()); + results.addAll(functions); + } + + return results.stream(); + } + + private String buildFunctionContext(StaticFunction func) { + var sb = new StringBuilder(); + sb.append(String.format("[FUNCTION] %s at %s\n", func.getName(), func.getAddress())); + + var callers = xrefRepository.findByCalleeId(func.getId()); + var callees = xrefRepository.findByCallerId(func.getId()); + + appendCallers(sb, callers); + appendCallees(sb, callees); + + sb.append("\n"); + sb.append("[ASSEMBLY]\n"); + sb.append(AiUtils.truncate(func.getAssembly(), functionContextMaxChars)); + + if (DEFAULT_CALL_GRAPH_HOPS > 0) { + appendNeighborContext(sb, callers, callees, func); + } + + return sb.toString(); + } + + private void appendCallers(StringBuilder sb, List callers) { + if (callers.isEmpty()) { + return; + } + + sb.append("Callers: "); + sb.append( + String.join(", ", callers.stream().limit(5).map(x -> x.getCaller().getName()).toList())); + sb.append("\n"); + } + + private void appendCallees(StringBuilder sb, List callees) { + if (callees.isEmpty()) { + return; + } + + sb.append("Callees: "); + sb.append( + String.join(", ", callees.stream().limit(10).map(x -> x.getCallee().getName()).toList())); + sb.append("\n"); + } + + private void appendNeighborContext( + StringBuilder sb, List callers, List callees, StaticFunction func) { + var neighbors = new ArrayList(); + neighbors.addAll(callers.stream().limit(2).map(StaticXref::getCaller).toList()); + neighbors.addAll(callees.stream().limit(2).map(StaticXref::getCallee).toList()); + + for (var neighbor : neighbors) { + if (neighbor.getId().equals(func.getId())) { + continue; + } + + sb.append("\n").append("[RELATED] ").append(neighbor.getName()); + sb.append(" at ").append(neighbor.getAddress()).append("\n"); + sb.append(AiUtils.truncate(neighbor.getAssembly(), NEIGHBOR_CONTEXT_MAX_CHARS)).append("\n"); + } + } + + private StaticAnalysis getLatestAnalysis(UUID binaryId) { + return analysisRepository.findFirstByBinaryIdOrderByCreatedAtDesc(binaryId).orElse(null); + } +} diff --git a/src/main/java/ai/decompile/ai/service/ChatMessageConverter.java b/src/main/java/ai/decompile/ai/service/ChatMessageConverter.java new file mode 100644 index 0000000..76cebeb --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/ChatMessageConverter.java @@ -0,0 +1,52 @@ +package ai.decompile.ai.service; + +import ai.decompile.ai.model.entity.ChatMessage; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +@Component +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class ChatMessageConverter { + + @Value("${app.ai.chat.max-history-messages:20}") + private int maxHistoryMessages; + + public List buildMessageList( + String systemPrompt, List historyMessages, String userMessageContent) { + var messages = new ArrayList(); + messages.add(new SystemMessage(systemPrompt)); + + for (var histMsg : historyMessages) { + if ("USER".equals(histMsg.getRole())) { + messages.add(new UserMessage(histMsg.getContent())); + } else { + messages.add(new AssistantMessage(histMsg.getContent())); + } + } + + messages.add(new UserMessage(userMessageContent)); + return messages; + } + + public List getRecentHistoryExcludingLastMessage(List messages) { + if (messages.isEmpty()) { + return Collections.emptyList(); + } + + var msgList = new ArrayList<>(messages); + msgList.removeLast(); + if (msgList.size() > maxHistoryMessages) { + return msgList.subList(msgList.size() - maxHistoryMessages, msgList.size()); + } + + return msgList; + } +} diff --git a/src/main/java/ai/decompile/ai/service/ChatService.java b/src/main/java/ai/decompile/ai/service/ChatService.java new file mode 100644 index 0000000..7baddba --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/ChatService.java @@ -0,0 +1,65 @@ +package ai.decompile.ai.service; + +import ai.decompile.ai.model.entity.ChatSession; +import ai.decompile.workspace.model.entity.Binary; +import ai.decompile.workspace.service.BinaryService; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; +import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; +import org.springframework.ai.chat.messages.Message; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Service; +import reactor.core.publisher.Flux; + +@Log4j2 +@Service +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class ChatService { + private final StreamingModelProvider streamingModelProvider; + private final ChatSessionService sessionService; + private final ChatMessageConverter messageConverter; + private final ChatContextBuilder contextBuilder; + private final BinaryService binaryService; + + public Flux chat(UUID sessionId, String userMessageContent) { + var session = sessionService.getSession(sessionId); + var binary = binaryService.getBinary(session.getBinaryId()); + + sessionService.saveMessage(session, "USER", userMessageContent); + + var messages = buildChatMessages(session, binary, userMessageContent); + var streamingModel = streamingModelProvider.resolve(); + var assistantContent = new AtomicReference<>(""); + + return streamingModel.stream(messages) + .doOnNext(token -> assistantContent.updateAndGet(current -> current + token)) + .doOnComplete(() -> saveAssistantMessage(session, assistantContent)) + .doOnError( + error -> log.error("Error during chat streaming for session {}", sessionId, error)); + } + + private Message[] buildChatMessages( + ChatSession session, Binary binary, String userMessageContent) { + var historyMessages = sessionService.getMessages(session.getId()); + + var context = contextBuilder.buildContext(session.getBinaryId(), userMessageContent, binary); + var systemPrompt = contextBuilder.buildSystemPrompt(context, binary); + + return messageConverter + .buildMessageList( + systemPrompt, + messageConverter.getRecentHistoryExcludingLastMessage(historyMessages), + userMessageContent) + .toArray(new Message[0]); + } + + private void saveAssistantMessage(ChatSession session, AtomicReference assistantContent) { + var fullContent = assistantContent.get(); + if (Objects.nonNull(fullContent) && !fullContent.isBlank()) { + sessionService.saveMessage(session, "ASSISTANT", fullContent); + } + } +} diff --git a/src/main/java/ai/decompile/ai/service/ChatSessionService.java b/src/main/java/ai/decompile/ai/service/ChatSessionService.java new file mode 100644 index 0000000..009721a --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/ChatSessionService.java @@ -0,0 +1,86 @@ +package ai.decompile.ai.service; + +import ai.decompile.ai.model.entity.ChatMessage; +import ai.decompile.ai.model.entity.ChatSession; +import ai.decompile.ai.model.repository.ChatMessageRepository; +import ai.decompile.ai.model.repository.ChatSessionRepository; +import ai.decompile.common.exception.NotFoundException; +import ai.decompile.workspace.service.BinaryService; +import java.util.List; +import java.util.Objects; +import java.util.UUID; +import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Propagation; +import org.springframework.transaction.annotation.Transactional; + +@Log4j2 +@Service +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class ChatSessionService { + private final ChatSessionRepository sessionRepository; + private final ChatMessageRepository messageRepository; + private final BinaryService binaryService; + private final StreamingModelProvider streamingModelProvider; + + @Value("${app.ai.chat.default-model:gemma4:12b}") + private String defaultModel; + + @Transactional + public ChatSession createSession(UUID binaryId, UUID projectId, String model) { + var binary = binaryService.getBinary(binaryId); + var effectiveModel = Objects.nonNull(model) ? model : defaultModel; + streamingModelProvider.resolve(); + + var session = + ChatSession.builder() + .binaryId(binaryId) + .projectId(projectId) + .title(binary.getFilename() + " — Chat") + .chatModel(effectiveModel) + .build(); + + return sessionRepository.save(session); + } + + @Transactional(readOnly = true) + public List getSessions(UUID binaryId) { + return sessionRepository.findByBinaryIdOrderByUpdatedAtDesc(binaryId); + } + + @Transactional(readOnly = true) + public ChatSession getSession(UUID sessionId) { + return sessionRepository + .findById(sessionId) + .orElseThrow( + () -> new NotFoundException("Chat session with id '" + sessionId + "' not found")); + } + + @Transactional + public void deleteSession(UUID sessionId) { + var session = getSession(sessionId); + sessionRepository.delete(session); + } + + @Transactional(readOnly = true) + public List getMessages(UUID sessionId) { + return messageRepository.findBySessionIdOrderByCreatedAtAsc(sessionId); + } + + @Transactional(propagation = Propagation.REQUIRES_NEW) + public void saveMessage(ChatSession session, String role, String content) { + var msg = + ChatMessage.builder() + .session(session) + .role(role) + .content(content) + .tokenCount(AiUtils.estimateTokens(content)) + .build(); + + messageRepository.save(msg); + } +} diff --git a/src/main/java/ai/decompile/ai/service/EmbeddingChunkBuilder.java b/src/main/java/ai/decompile/ai/service/EmbeddingChunkBuilder.java new file mode 100644 index 0000000..3f06cf7 --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/EmbeddingChunkBuilder.java @@ -0,0 +1,179 @@ +package ai.decompile.ai.service; + +import ai.decompile.ai.model.entity.EmbeddingChunk; +import ai.decompile.analysis.model.entity.StaticFunction; +import ai.decompile.analysis.model.entity.StaticXref; +import ai.decompile.analysis.model.repository.StaticLabelRepository; +import ai.decompile.analysis.model.repository.StaticXrefRepository; +import ai.decompile.workspace.model.entity.Binary; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; +import java.util.function.Function; +import lombok.RequiredArgsConstructor; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +@Component +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class EmbeddingChunkBuilder { + private final StaticXrefRepository xrefRepository; + private final StaticLabelRepository labelRepository; + private final ObjectMapper objectMapper; + + @Value("${app.ai.embedding.chunk-max-chars:3000}") + private int chunkMaxChars; + + public List buildFunctionChunks(List functions, Binary binary) { + var graphContext = loadGraphContext(functions); + + var chunks = new ArrayList(); + for (var func : functions) { + chunks.add( + EmbeddingChunk.builder() + .id(UUID.randomUUID()) + .binaryId(binary.getId()) + .chunkType("FUNCTION") + .sourceType("STATIC_FUNCTION") + .sourceId(func.getId()) + .content(buildFunctionChunkContent(func, binary, graphContext)) + .metadata(AiUtils.toJson(objectMapper, buildFunctionMetadata(func, graphContext))) + .tokenCount(AiUtils.estimateTokens(func.getAssembly())) + .createdAt(java.time.Instant.now()) + .build()); + } + + return chunks; + } + + public EmbeddingChunk buildMetadataChunk(Binary binary, int functionCount) { + var content = + String.format( + "[METADATA] Binary Overview\nFilename: %s\nFormat: %s\nArchitecture: %s\nCompiler: %s\nFunction count: %d\n", + AiUtils.orUnknown(binary.getFilename()), + AiUtils.orUnknown(binary.getFormat()), + AiUtils.orUnknown(binary.getArchitecture()), + AiUtils.orUnknown(binary.getCompiler()), + functionCount); + + return EmbeddingChunk.builder() + .id(UUID.randomUUID()) + .binaryId(binary.getId()) + .chunkType("METADATA") + .content(content) + .metadata( + AiUtils.toJson( + objectMapper, Map.of("chunk_type", "METADATA", "function_count", functionCount))) + .tokenCount(AiUtils.estimateTokens(content)) + .createdAt(java.time.Instant.now()) + .build(); + } + + private FunctionGraphContext loadGraphContext(List functions) { + var nameToCallers = + loadXrefNames(functions, xrefRepository::findByCalleeId, xr -> xr.getCaller().getName()); + + var nameToCallees = + loadXrefNames(functions, xrefRepository::findByCallerId, xr -> xr.getCallee().getName()); + + var funcToLabels = loadLabels(functions); + + return new FunctionGraphContext(nameToCallers, nameToCallees, funcToLabels); + } + + private Map> loadXrefNames( + List functions, + Function> queryFn, + Function nameFn) { + var m = new HashMap>(); + for (var f : functions) { + var xrefs = queryFn.apply(f.getId()); + if (!xrefs.isEmpty()) { + m.put(f.getName(), xrefs.stream().map(nameFn).distinct().toList()); + } + } + + return m; + } + + private Map> loadLabels(List functions) { + var m = new HashMap>(); + for (var f : functions) { + var labels = labelRepository.findByFunctionId(f.getId()); + if (!labels.isEmpty()) { + m.put( + f.getId(), + labels.stream().map(l -> l.getAddress() + " (" + l.getName() + ")").toList()); + } + } + + return m; + } + + private String buildFunctionChunkContent( + StaticFunction func, Binary binary, FunctionGraphContext ctx) { + var sb = new StringBuilder(); + sb.append(String.format("[FUNCTION] %s at %s\n", func.getName(), func.getAddress())); + sb.append( + String.format( + "Binary: %s | %s | %s | %s\n", + AiUtils.orUnknown(binary.getFilename()), + AiUtils.orUnknown(binary.getArchitecture()), + AiUtils.orUnknown(binary.getCompiler()), + AiUtils.orUnknown(binary.getFormat()))); + + var callers = ctx.nameToCallers.getOrDefault(func.getName(), List.of()); + if (!callers.isEmpty()) { + sb.append("Callers: ").append(String.join(", ", callers)).append("\n"); + } + + var callees = ctx.nameToCallees.getOrDefault(func.getName(), List.of()); + if (!callees.isEmpty()) { + sb.append("Callees: ").append(String.join(", ", callees)).append("\n"); + } + + var labels = ctx.funcToLabels.getOrDefault(func.getId(), List.of()); + if (!labels.isEmpty()) { + sb.append("Labels:\n"); + for (var l : labels) { + sb.append(" ").append(l).append("\n"); + } + } + + sb.append("\n"); + if (Objects.nonNull(func.getDecompiledCode()) && !func.getDecompiledCode().isBlank()) { + sb.append("[DECOMPILED]\n") + .append(AiUtils.truncate(func.getDecompiledCode(), chunkMaxChars / 2)) + .append("\n"); + } + + sb.append("[ASSEMBLY]\n").append(AiUtils.truncate(func.getAssembly(), chunkMaxChars)); + + return sb.toString(); + } + + private Map buildFunctionMetadata(StaticFunction func, FunctionGraphContext ctx) { + Map m = new HashMap<>(); + m.put("function_name", func.getName()); + m.put("address", func.getAddress()); + m.put("callers", ctx.nameToCallers.getOrDefault(func.getName(), List.of())); + m.put("callees", ctx.nameToCallees.getOrDefault(func.getName(), List.of())); + m.put("label_count", ctx.funcToLabels.getOrDefault(func.getId(), List.of()).size()); + m.put("assembly_length", Objects.nonNull(func.getAssembly()) ? func.getAssembly().length() : 0); + m.put( + "has_decompiled_code", + Objects.nonNull(func.getDecompiledCode()) && !func.getDecompiledCode().isBlank()); + return m; + } + + record FunctionGraphContext( + Map> nameToCallers, + Map> nameToCallees, + Map> funcToLabels) {} +} diff --git a/src/main/java/ai/decompile/ai/service/EmbeddingClient.java b/src/main/java/ai/decompile/ai/service/EmbeddingClient.java new file mode 100644 index 0000000..913e244 --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/EmbeddingClient.java @@ -0,0 +1,9 @@ +package ai.decompile.ai.service; + +import java.util.List; + +public interface EmbeddingClient { + float[] embed(String text); + + List embed(List texts); +} diff --git a/src/main/java/ai/decompile/ai/service/EmbeddingService.java b/src/main/java/ai/decompile/ai/service/EmbeddingService.java new file mode 100644 index 0000000..94c081c --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/EmbeddingService.java @@ -0,0 +1,97 @@ +package ai.decompile.ai.service; + +import ai.decompile.ai.model.entity.EmbeddingChunk; +import ai.decompile.ai.model.repository.EmbeddingChunkRepository; +import ai.decompile.analysis.model.entity.StaticAnalysis; +import ai.decompile.analysis.model.repository.StaticAnalysisRepository; +import ai.decompile.workspace.service.BinaryService; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.UUID; +import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; + +@Log4j2 +@Service +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class EmbeddingService { + private final EmbeddingClient embeddingClient; + private final EmbeddingChunkBuilder chunkBuilder; + private final EmbeddingChunkRepository embeddingChunkRepository; + private final StaticAnalysisRepository analysisRepository; + private final BinaryService binaryService; + + @Value("${app.ai.embedding.batch-size:20}") + private int batchSize; + + @Value("${app.ai.embedding.similarity-threshold:0.6}") + private double similarityThreshold; + + public float[] embed(String text) { + return embeddingClient.embed(text); + } + + @Transactional + public int indexBinary(UUID binaryId) { + log.info("Starting embedding generation for binary {}", binaryId); + embeddingChunkRepository.deleteByBinaryId(binaryId); + embeddingChunkRepository.flush(); + + var binary = binaryService.getBinary(binaryId); + var analysis = getLatestAnalysis(binaryId); + if (Objects.isNull(analysis)) { + throw new IllegalStateException( + "No static analysis found for binary " + binaryId + ". Run analysis first."); + } + + var functions = analysis.getFunctions(); + log.info("Indexing {} functions for binary {}", functions.size(), binaryId); + + var functionChunks = chunkBuilder.buildFunctionChunks(functions, binary); + var allChunks = new ArrayList<>(functionChunks); + allChunks.add(chunkBuilder.buildMetadataChunk(binary, functions.size())); + embedAndPersistChunks(allChunks); + + log.info("Embedding complete for binary {}: {} chunks", binaryId, allChunks.size()); + return allChunks.size(); + } + + private void embedAndPersistChunks(List chunks) { + for (var i = 0; i < chunks.size(); i += batchSize) { + var batch = chunks.subList(i, Math.min(i + batchSize, chunks.size())); + var texts = batch.stream().map(EmbeddingChunk::getContent).toList(); + var embeddings = embeddingClient.embed(texts); + + for (int j = 0; j < batch.size(); j++) { + var c = batch.get(j); + embeddingChunkRepository.insertChunk( + c.getId(), + c.getBinaryId(), + c.getChunkType(), + c.getSourceType(), + c.getSourceId(), + c.getContent(), + AiUtils.toVectorString(embeddings.get(j)), + c.getMetadata(), + c.getTokenCount(), + c.getCreatedAt()); + } + + log.debug( + "Saved batch {}-{}/{} chunks", + i + 1, + Math.min(i + batchSize, chunks.size()), + chunks.size()); + } + } + + private StaticAnalysis getLatestAnalysis(UUID binaryId) { + return analysisRepository.findFirstByBinaryIdOrderByCreatedAtDesc(binaryId).orElse(null); + } +} diff --git a/src/main/java/ai/decompile/ai/service/OllamaEmbeddingClient.java b/src/main/java/ai/decompile/ai/service/OllamaEmbeddingClient.java new file mode 100644 index 0000000..e8f9a6a --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/OllamaEmbeddingClient.java @@ -0,0 +1,28 @@ +package ai.decompile.ai.service; + +import java.util.List; +import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; +import org.springframework.ai.ollama.OllamaEmbeddingModel; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.resilience.annotation.Retryable; +import org.springframework.stereotype.Component; + +@Log4j2 +@Component +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class OllamaEmbeddingClient implements EmbeddingClient { + private final OllamaEmbeddingModel embeddingModel; + + @Override + @Retryable(maxRetries = 2, delay = 500, multiplier = 2) + public float[] embed(String text) { + return embeddingModel.embed(text); + } + + @Override + public List embed(List texts) { + return embeddingModel.embed(texts); + } +} diff --git a/src/main/java/ai/decompile/ai/service/OllamaStreamingModelProvider.java b/src/main/java/ai/decompile/ai/service/OllamaStreamingModelProvider.java new file mode 100644 index 0000000..59d8ccf --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/OllamaStreamingModelProvider.java @@ -0,0 +1,29 @@ +package ai.decompile.ai.service; + +import java.util.Objects; +import lombok.RequiredArgsConstructor; +import lombok.extern.log4j.Log4j2; +import org.springframework.ai.chat.model.StreamingChatModel; +import org.springframework.ai.ollama.OllamaChatModel; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +@Log4j2 +@Component +@RequiredArgsConstructor +@ConditionalOnProperty(name = "app.ai.enabled", havingValue = "true", matchIfMissing = true) +public class OllamaStreamingModelProvider implements StreamingModelProvider { + private final ObjectProvider ollamaChatModelProvider; + + @Override + public StreamingChatModel resolve() { + var ollamaModel = ollamaChatModelProvider.getIfAvailable(); + if (Objects.isNull(ollamaModel)) { + throw new IllegalStateException( + "Ollama is not available. Install ollama and run 'ollama pull gemma4:12b'."); + } + + return ollamaModel; + } +} diff --git a/src/main/java/ai/decompile/ai/service/StreamingModelProvider.java b/src/main/java/ai/decompile/ai/service/StreamingModelProvider.java new file mode 100644 index 0000000..e19a533 --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/StreamingModelProvider.java @@ -0,0 +1,8 @@ +package ai.decompile.ai.service; + +import org.springframework.ai.chat.model.StreamingChatModel; + +@FunctionalInterface +public interface StreamingModelProvider { + StreamingChatModel resolve(); +} diff --git a/src/main/java/ai/decompile/ai/service/package-info.java b/src/main/java/ai/decompile/ai/service/package-info.java new file mode 100644 index 0000000..065b58c --- /dev/null +++ b/src/main/java/ai/decompile/ai/service/package-info.java @@ -0,0 +1 @@ +package ai.decompile.ai.service; diff --git a/src/main/resources/application.yml b/src/main/resources/application.yml index d30a2be..eca8b23 100644 --- a/src/main/resources/application.yml +++ b/src/main/resources/application.yml @@ -4,6 +4,8 @@ spring: enabled: true application: name: decompile-ai + config: + import: optional:file:.env[.properties] datasource: url: jdbc:postgresql://localhost:5432/decompile_ai username: decompile_ai @@ -22,6 +24,18 @@ spring: multipart: max-file-size: 500MB max-request-size: 500MB + ai: + ollama: + base-url: http://localhost:11434 + embedding: + options: + model: ${AI_EMBEDDING_MODEL:qwen3-embedding:latest} + chat: + options: + model: ${AI_CHAT_MODEL_OLLAMA:gemma4:12b} + temperature: 0.3 + top-p: 0.95 + top-k: 64 app: storage: @@ -43,7 +57,24 @@ app: ghidra: image: decompile-ai/ghidra:latest timeout-seconds: 600 + ai: + default-chat-provider: ${AI_DEFAULT_CHAT_PROVIDER:ollama} + embedding: + dimension: ${AI_EMBEDDING_DIMENSION:4096} + chunk-max-chars: 3000 + batch-size: 20 + rag: + top-k: 10 + similarity-threshold: 0.6 + max-context-tokens: 8000 + fallback-min-ratio: 0.4 + function-context-max-chars: 5000 + chat: + max-history-messages: 20 + temperature: 0.3 logging: level: org.springframework.modulith: DEBUG + ai.decompile.ai: DEBUG + org.springframework.ai: INFO diff --git a/src/main/resources/db/migration/V008__ai_embedding_and_chat.sql b/src/main/resources/db/migration/V008__ai_embedding_and_chat.sql new file mode 100644 index 0000000..0b30314 --- /dev/null +++ b/src/main/resources/db/migration/V008__ai_embedding_and_chat.sql @@ -0,0 +1,43 @@ +CREATE EXTENSION IF NOT EXISTS vector; + +CREATE TABLE embedding_chunk +( + id UUID PRIMARY KEY, + binary_id UUID NOT NULL REFERENCES binaries (id) ON DELETE CASCADE, + chunk_type VARCHAR(30) NOT NULL, + source_type VARCHAR(30), + source_id UUID, + content TEXT NOT NULL, + embedding vector(4096), + metadata JSONB, + token_count INTEGER, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_embedding_chunk_binary ON embedding_chunk (binary_id); +CREATE INDEX idx_embedding_chunk_source ON embedding_chunk (source_type, source_id); + +CREATE TABLE chat_session +( + id UUID PRIMARY KEY, + binary_id UUID NOT NULL REFERENCES binaries (id) ON DELETE CASCADE, + project_id UUID NOT NULL REFERENCES projects (id) ON DELETE CASCADE, + title VARCHAR(255), + chat_model VARCHAR(100) NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_chat_session_binary ON chat_session (binary_id); + +CREATE TABLE chat_message +( + id UUID PRIMARY KEY, + session_id UUID NOT NULL REFERENCES chat_session (id) ON DELETE CASCADE, + role VARCHAR(20) NOT NULL, + content TEXT NOT NULL, + token_count INTEGER, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_chat_message_session ON chat_message (session_id, created_at);