Decompile.AI/backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai): add AI chat and RAG embedding module

Browse Source 
Core AI module with chat sessions, message history, pgvector embeddings,
semantic search via RAG, and Ollama integration for chat/embedding models.

Includes:
- Spring AI Ollama dependency (BOM 1.0.1)
- Entity model: ChatSession, ChatMessage, EmbeddingChunk
- Custom PostgreSQLVectorJdbcType for pgvector mapping
- Services: ChatService, EmbeddingService, ChatContextBuilder
- ChatController with SSE streaming endpoint
- Flyway migration V008 (pgvector extension + tables)
- AI configuration in application.yml
This commit is contained in:
Rodrigo Verdiani 2026-06-08 20:12:59 -03:00
parent 199345d497
commit d01ed613fb
33 changed files with 1641 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
.env.example Normal file
View File

@ -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

63
pom.xml
View File

@ -26,10 +26,11 @@
<tag/>
<url/>
</scm>
<properties>
<java.version>25</java.version>
<spring-modulith.version>2.0.6</spring-modulith.version>
</properties>
<properties>
<java.version>25</java.version>
<spring-modulith.version>2.0.6</spring-modulith.version>
<spring-ai.version>1.0.1</spring-ai.version>
</properties>
<dependencies>
<dependency>
<groupId>org.springframework.boot</groupId>
@ -67,7 +68,6 @@
<dependency>
<groupId>org.postgresql</groupId>
<artifactId>postgresql</artifactId>
<scope>runtime</scope>
</dependency>
<dependency>
<groupId>org.springframework.modulith</groupId>
@ -135,25 +135,42 @@
<artifactId>docker-java-transport-httpclient5</artifactId>
<version>3.7.1</version>
</dependency>
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
<version>3.20.0</version>
<scope>compile</scope>
</dependency>
</dependencies>
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
<version>3.20.0</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-ollama</artifactId>
<exclusions>
<exclusion>
<groupId>io.swagger.core.v3</groupId>
<artifactId>swagger-annotations-jakarta</artifactId>
</exclusion>
</exclusions>
</dependency>
</dependencies>
<dependencyManagement>
<dependencies>
<dependency>
<groupId>org.springframework.modulith</groupId>
<artifactId>spring-modulith-bom</artifactId>
<version>${spring-modulith.version}</version>
<type>pom</type>
<scope>import</scope>
</dependency>
</dependencies>
</dependencyManagement>
<dependencyManagement>
<dependencies>
<dependency>
<groupId>org.springframework.modulith</groupId>
<artifactId>spring-modulith-bom</artifactId>
<version>${spring-modulith.version}</version>
<type>pom</type>
<scope>import</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-bom</artifactId>
<version>${spring-ai.version}</version>
<type>pom</type>
<scope>import</scope>
</dependency>
</dependencies>
</dependencyManagement>
<build>
<plugins>

61
src/main/java/ai/decompile/ai/config/AiConfig.java Normal file
View File

@ -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();
}
}

115
src/main/java/ai/decompile/ai/controller/ChatController.java Normal file
View File

@ -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<SessionResponse> 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<List<SessionResponse>> 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<SessionResponse> 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<Void> 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<List<MessageResponse>> 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<ServerSentEvent<String>> 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<String> event(String type, String key, Object value) {
return ServerSentEvent.<String>builder().data(toJson(Map.of("type", type, key, value))).build();
}
private ServerSentEvent<String> errorEvent(Throwable e) {
String msg = e.getMessage() != null ? e.getMessage() : "Unknown error";
return ServerSentEvent.<String>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\"}";
}
}
}

1
src/main/java/ai/decompile/ai/controller/package-info.java Normal file
View File

@ -0,0 +1 @@
package ai.decompile.ai.controller;

6
src/main/java/ai/decompile/ai/model/dto/ChatRequest.java Normal file
View File

@ -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) {}

18
src/main/java/ai/decompile/ai/model/dto/MessageResponse.java Normal file
View File

@ -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());
}
}

27
src/main/java/ai/decompile/ai/model/dto/SessionResponse.java Normal file
View File

@ -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());
}
}

1
src/main/java/ai/decompile/ai/model/dto/package-info.java Normal file
View File

@ -0,0 +1 @@
package ai.decompile.ai.model.dto;

37
src/main/java/ai/decompile/ai/model/entity/ChatMessage.java Normal file
View File

@ -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;
}

60
src/main/java/ai/decompile/ai/model/entity/ChatSession.java Normal file
View File

@ -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<ChatMessage> 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;
}

55
src/main/java/ai/decompile/ai/model/entity/EmbeddingChunk.java Normal file
View File

@ -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;
}

1
src/main/java/ai/decompile/ai/model/entity/package-info.java Normal file
View File

@ -0,0 +1 @@
package ai.decompile.ai.model.entity;

115
src/main/java/ai/decompile/ai/model/jdbc/PostgreSQLVectorJdbcType.java Normal file
View File

@ -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 <X> ValueBinder<X> getBinder(JavaType<X> 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 <X> ValueExtractor<X> getExtractor(JavaType<X> 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();
}
}

10
src/main/java/ai/decompile/ai/model/repository/ChatMessageRepository.java Normal file
View File

@ -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<ChatMessage, UUID> {
List<ChatMessage> findBySessionIdOrderByCreatedAtAsc(UUID sessionId);
}

10
src/main/java/ai/decompile/ai/model/repository/ChatSessionRepository.java Normal file
View File

@ -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<ChatSession, UUID> {
List<ChatSession> findByBinaryIdOrderByUpdatedAtDesc(UUID binaryId);
}

65
src/main/java/ai/decompile/ai/model/repository/EmbeddingChunkRepository.java Normal file
View File

@ -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<EmbeddingChunk, UUID> {
@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<EmbeddingChunk> 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<EmbeddingChunk> searchByContent(
@Param("binaryId") UUID binaryId, @Param("term") String term, @Param("limit") int limit);
}

1
src/main/java/ai/decompile/ai/model/repository/package-info.java Normal file
View File

@ -0,0 +1 @@
package ai.decompile.ai.model.repository;

10
src/main/java/ai/decompile/ai/package-info.java Normal file
View File

@ -0,0 +1,10 @@
@org.springframework.modulith.ApplicationModule(
displayName = "AI",
allowedDependencies = {
"workspace::entities",
"workspace::services",
"analysis::entities",
"analysis::services",
"common"
})
package ai.decompile.ai;

50
src/main/java/ai/decompile/ai/service/AiUtils.java Normal file
View File

@ -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 "{}";
}
}
}

309
src/main/java/ai/decompile/ai/service/ChatContextBuilder.java Normal file
View File

@ -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<EmbeddingChunk>();
var inlineSections = new ArrayList<String>();
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<EmbeddingChunk> 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<EmbeddingChunk> merged, UUID binaryId, List<String> 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<EmbeddingChunk> 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<String> inlineSections,
LinkedHashSet<EmbeddingChunk> 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<EmbeddingChunk> merged, StaticFunction func) {
return merged.stream()
.anyMatch(c -> Objects.nonNull(c.getSourceId()) && c.getSourceId().equals(func.getId()));
}
private String formatContext(
List<String> inlineSections, LinkedHashSet<EmbeddingChunk> 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<EmbeddingChunk> semanticSearch(UUID binaryId, String query) {
var queryEmbedding = embeddingService.embed(query);
var vectorStr = AiUtils.toVectorString(queryEmbedding);
return embeddingChunkRepository.findSimilar(binaryId, vectorStr, similarityThreshold, topK);
}
private List<String> extractMeaningfulTokens(String query) {
return Stream.of(query.split("[\\s,.;:!?()\\[\\]{}\"']+"))
.filter(t -> t.length() >= MIN_TOKEN_LENGTH)
.distinct()
.limit(10)
.toList();
}
private Stream<StaticFunction> 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<String>();
while (matcher.find() && extraTokens.size() < 10) {
extraTokens.add(matcher.group(1));
}
var allTokens = Stream.concat(tokens.stream(), extraTokens.stream()).distinct().toList();
Set<StaticFunction> 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<StaticXref> 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<StaticXref> 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<StaticXref> callers, List<StaticXref> callees, StaticFunction func) {
var neighbors = new ArrayList<StaticFunction>();
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);
}
}

52
src/main/java/ai/decompile/ai/service/ChatMessageConverter.java Normal file
View File

@ -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<Message> buildMessageList(
String systemPrompt, List<ChatMessage> historyMessages, String userMessageContent) {
var messages = new ArrayList<Message>();
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<ChatMessage> getRecentHistoryExcludingLastMessage(List<ChatMessage> 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;
}
}

65
src/main/java/ai/decompile/ai/service/ChatService.java Normal file
View File

@ -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<String> 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<String> assistantContent) {
var fullContent = assistantContent.get();
if (Objects.nonNull(fullContent) && !fullContent.isBlank()) {
sessionService.saveMessage(session, "ASSISTANT", fullContent);
}
}
}

86
src/main/java/ai/decompile/ai/service/ChatSessionService.java Normal file
View File

@ -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<ChatSession> 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<ChatMessage> 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);
}
}

179
src/main/java/ai/decompile/ai/service/EmbeddingChunkBuilder.java Normal file
View File

@ -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<EmbeddingChunk> buildFunctionChunks(List<StaticFunction> functions, Binary binary) {
var graphContext = loadGraphContext(functions);
var chunks = new ArrayList<EmbeddingChunk>();
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<StaticFunction> 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<String, List<String>> loadXrefNames(
List<StaticFunction> functions,
Function<UUID, List<StaticXref>> queryFn,
Function<StaticXref, String> nameFn) {
var m = new HashMap<String, List<String>>();
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<UUID, List<String>> loadLabels(List<StaticFunction> functions) {
var m = new HashMap<UUID, List<String>>();
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<String, Object> buildFunctionMetadata(StaticFunction func, FunctionGraphContext ctx) {
Map<String, Object> 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<String, List<String>> nameToCallers,
Map<String, List<String>> nameToCallees,
Map<UUID, List<String>> funcToLabels) {}
}

9
src/main/java/ai/decompile/ai/service/EmbeddingClient.java Normal file
View File

@ -0,0 +1,9 @@
package ai.decompile.ai.service;
import java.util.List;
public interface EmbeddingClient {
float[] embed(String text);
List<float[]> embed(List<String> texts);
}

97
src/main/java/ai/decompile/ai/service/EmbeddingService.java Normal file
View File

@ -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<EmbeddingChunk> 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);
}
}

28
src/main/java/ai/decompile/ai/service/OllamaEmbeddingClient.java Normal file
View File

@ -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<float[]> embed(List<String> texts) {
return embeddingModel.embed(texts);
}
}

29
src/main/java/ai/decompile/ai/service/OllamaStreamingModelProvider.java Normal file
View File

@ -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<OllamaChatModel> 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;
}
}

8
src/main/java/ai/decompile/ai/service/StreamingModelProvider.java Normal file
View File

@ -0,0 +1,8 @@
package ai.decompile.ai.service;
import org.springframework.ai.chat.model.StreamingChatModel;
@FunctionalInterface
public interface StreamingModelProvider {
StreamingChatModel resolve();
}

1
src/main/java/ai/decompile/ai/service/package-info.java Normal file
View File

@ -0,0 +1 @@
package ai.decompile.ai.service;

31
src/main/resources/application.yml
View File

@ -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

43
src/main/resources/db/migration/V008__ai_embedding_and_chat.sql Normal file
View File

@ -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);